models/moa_model.py
def forward_rnn(self, input_dict, state, seq_lens):
        """
        Forward pass through the MOA LSTMs.
        Implicitly assigns the value function output to self_value_out, and does not return this.
        :param input_dict: The input tensors.
        :param state: The model state.
        :param seq_lens: LSTM sequence lengths.
        :return: The policy logits and new LSTM states.
        """
        # Evaluate the actor-critic model
        pass_dict = {"curr_obs": input_dict["ac_trunk"]}
        h1, c1, h2, c2, *_ = state
        (self._model_out, self._value_out, output_h1, output_c1,) = self.actions_model.forward_rnn(
            pass_dict, [h1, c1], seq_lens
        )

        # Evaluate the MOA, and generate counterfactual actions.
        # To do this: cycle through all possible actions and get predictions for what other
        # agents would do if this action was taken at each trajectory step.

        # First we have to evaluate the MOA over the true trajectory to obtain the hidden state
        # that we will use for the next timestep's counterfactual predictions.
        prev_moa_trunk = input_dict["prev_moa_trunk"]
        other_actions = input_dict["other_agent_actions"]
        agent_action = tf.expand_dims(input_dict["prev_actions"], axis=-1)
        all_actions = tf.concat([agent_action, other_actions], axis=-1, name="concat_true_actions")
        self._true_one_hot_actions = self._reshaped_one_hot_actions(all_actions, "forward_one_hot")
        true_action_pass_dict = {
            "curr_obs": prev_moa_trunk,
            "prev_total_actions": self._true_one_hot_actions,
        }

        # Compute the true action prediction, used to determine the MOA loss.
        self._action_pred, output_h2, output_c2 = self.moa_model.forward_rnn(
            true_action_pass_dict, [h2, c2], seq_lens
        )

        # Make counterfactual predictions, used for computing the influence reward.
        counterfactual_preds = []
        for i in range(self.num_outputs):
            # Shape of other_actions is (num_envs, ?, num_other_agents)
            # To add the counterfactual action to it, other_actions can be padded with the constant
            # action value.
            actions_with_counterfactual = tf.pad(
                other_actions, paddings=[[0, 0], [0, 0], [1, 0]], mode="CONSTANT", constant_values=i
            )
            one_hot_actions = self._reshaped_one_hot_actions(
                actions_with_counterfactual, "actions_with_counterfactual_one_hot"
            )
            pass_dict = {"curr_obs": prev_moa_trunk, "prev_total_actions": one_hot_actions}
            counterfactual_pred, _, _ = self.moa_model.forward_rnn(pass_dict, [h2, c2], seq_lens)
            counterfactual_preds.append(tf.expand_dims(counterfactual_pred, axis=-2))
        self._counterfactuals = tf.concat(
            counterfactual_preds, axis=-2, name="concat_counterfactuals"
        )

        # TODO(@evinitsky) move this into ppo_moa by using restore_original_dimensions()
        self._other_agent_actions = input_dict["other_agent_actions"]
        self._visibility = input_dict["visible_agents"]

        return self._model_out, [output_h1, output_c1, output_h2, output_c2]

        Forward pass through the MOA LSTMs.
        Implicitly assigns the value function output to self_value_out, and does not return this.
        :param input_dict: The input tensors.
        :param state: The model state.
        :param seq_lens: LSTM sequence lengths.
        :return: The policy logits and new LSTM states.
        
Forward pass through the MOA LSTMs.
Implicitly assigns the value function output to self_value_out, and does not return this.
Forward pass through the MOA LSTMs . Implicitly assigns the value function output to self_value_out and does not return this .
def forward_rnn(self, input_dict, state, seq_lens):
        pass_dict = {"curr_obs": input_dict["ac_trunk"]}
        h1, c1, h2, c2, *_ = state
        (self._model_out, self._value_out, output_h1, output_c1,) = self.actions_model.forward_rnn(
            pass_dict, [h1, c1], seq_lens
        )
        prev_moa_trunk = input_dict["prev_moa_trunk"]
        other_actions = input_dict["other_agent_actions"]
        agent_action = tf.expand_dims(input_dict["prev_actions"], axis=-1)
        all_actions = tf.concat([agent_action, other_actions], axis=-1, name="concat_true_actions")
        self._true_one_hot_actions = self._reshaped_one_hot_actions(all_actions, "forward_one_hot")
        true_action_pass_dict = {
            "curr_obs": prev_moa_trunk,
            "prev_total_actions": self._true_one_hot_actions,
        }
        self._action_pred, output_h2, output_c2 = self.moa_model.forward_rnn(
            true_action_pass_dict, [h2, c2], seq_lens
        )
        counterfactual_preds = []
        for i in range(self.num_outputs):
            actions_with_counterfactual = tf.pad(
                other_actions, paddings=[[0, 0], [0, 0], [1, 0]], mode="CONSTANT", constant_values=i
            )
            one_hot_actions = self._reshaped_one_hot_actions(
                actions_with_counterfactual, "actions_with_counterfactual_one_hot"
            )
            pass_dict = {"curr_obs": prev_moa_trunk, "prev_total_actions": one_hot_actions}
            counterfactual_pred, _, _ = self.moa_model.forward_rnn(pass_dict, [h2, c2], seq_lens)
            counterfactual_preds.append(tf.expand_dims(counterfactual_pred, axis=-2))
        self._counterfactuals = tf.concat(
            counterfactual_preds, axis=-2, name="concat_counterfactuals"
        )
        self._other_agent_actions = input_dict["other_agent_actions"]
        self._visibility = input_dict["visible_agents"]
        return self._model_out, [output_h1, output_c1, output_h2, output_c2]
def forward_rnn ( self , input_dict , state , seq_lens ) : pass_dict = { "curr_obs" : input_dict [ "ac_trunk" ] } h1 , c1 , h2 , c2 , * _ = state ( self . _model_out , self . _value_out , output_h1 , output_c1 , ) = self . actions_model . forward_rnn ( pass_dict , [ h1 , c1 ] , seq_lens ) prev_moa_trunk = input_dict [ "prev_moa_trunk" ] other_actions = input_dict [ "other_agent_actions" ] agent_action = tf . expand_dims ( input_dict [ "prev_actions" ] , axis = - 1 ) all_actions = tf . concat ( [ agent_action , other_actions ] , axis = - 1 , name = "concat_true_actions" ) self . _true_one_hot_actions = self . _reshaped_one_hot_actions ( all_actions , "forward_one_hot" ) true_action_pass_dict = { "curr_obs" : prev_moa_trunk , "prev_total_actions" : self . _true_one_hot_actions , } self . _action_pred , output_h2 , output_c2 = self . moa_model . forward_rnn ( true_action_pass_dict , [ h2 , c2 ] , seq_lens ) counterfactual_preds = [ ] for i in range ( self . num_outputs ) : actions_with_counterfactual = tf . pad ( other_actions , paddings = [ [ 0 , 0 ] , [ 0 , 0 ] , [ 1 , 0 ] ] , mode = "CONSTANT" , constant_values = i ) one_hot_actions = self . _reshaped_one_hot_actions ( actions_with_counterfactual , "actions_with_counterfactual_one_hot" ) pass_dict = { "curr_obs" : prev_moa_trunk , "prev_total_actions" : one_hot_actions } counterfactual_pred , _ , _ = self . moa_model . forward_rnn ( pass_dict , [ h2 , c2 ] , seq_lens ) counterfactual_preds . append ( tf . expand_dims ( counterfactual_pred , axis = - 2 ) ) self . _counterfactuals = tf . concat ( counterfactual_preds , axis = - 2 , name = "concat_counterfactuals" ) self . _other_agent_actions = input_dict [ "other_agent_actions" ] self . _visibility = input_dict [ "visible_agents" ] return self . _model_out , [ output_h1 , output_c1 , output_h2 , output_c2 ]
Forward pass through the MOA LSTMs.
Forward pass through the MOA LSTMs .
"""
        Forward pass through the MOA LSTMs.
        Implicitly assigns the value function output to self_value_out, and does not return this.
        :param input_dict: The input tensors.
        :param state: The model state.
        :param seq_lens: LSTM sequence lengths.
        :return: The policy logits and new LSTM states.
        """ # Evaluate the actor-critic model # Evaluate the MOA, and generate counterfactual actions. # To do this: cycle through all possible actions and get predictions for what other # agents would do if this action was taken at each trajectory step. # First we have to evaluate the MOA over the true trajectory to obtain the hidden state # that we will use for the next timestep's counterfactual predictions. # Compute the true action prediction, used to determine the MOA loss. # Make counterfactual predictions, used for computing the influence reward. # Shape of other_actions is (num_envs, ?, num_other_agents) # To add the counterfactual action to it, other_actions can be padded with the constant # action value. # TODO(@evinitsky) move this into ppo_moa by using restore_original_dimensions()
The policy logits and new LSTM states.
The policy logits and new LSTM states .
self
input_dict
The input tensors.
The input tensors .
state
The model state.
The model state .
seq_lens
LSTM sequence lengths.
LSTM sequence lengths .