Minor changes to GAIL

ml-agents/mlagents/trainers/ppo/reward_signals/curiosity/model.py

@@ -41,7 +41,7 @@ def create_curiosity_encoders(self):
                     self.policy_model.curiosity_enc_size,
                     LearningModel.swish, 1,
                     "stream_{}_visual_obs_encoder"
-                    .format(i), False)
+                        .format(i), False)
 
                 encoded_next_visual = self.policy_model.create_visual_obs_encoder(
                     self.next_visual_in[i],

ml-agents/mlagents/trainers/ppo/reward_signals/gail/model.py

@@ -1,5 +1,4 @@
 import tensorflow as tf
-import numpy as np
 
 
 class GAILModel(object):
@@ -14,6 +13,11 @@ def __init__(self, policy_model, h_size, lr):
     def make_inputs(self):
         self.obs_in_expert = tf.placeholder(
             shape=[None, self.policy_model.vec_obs_size], dtype=tf.float32)
+        self.done_expert = tf.placeholder(
+            shape=[None, 1], dtype=tf.float32)
+        self.done_policy =  tf.placeholder(
+            shape=[None, 1], dtype=tf.float32)
+
         if self.policy_model.brain.vector_action_space_type == 'continuous':
             action_length = self.policy_model.act_size[0]
             self.action_in_expert = tf.placeholder(
@@ -27,9 +31,9 @@ def make_inputs(self):
             tf.one_hot(self.action_in_expert[:, i], self.policy_model.act_size[i]) for i in
             range(len(self.policy_model.act_size))], axis=1)
 
-    def create_encoder(self, state_in, action_in, reuse):
+    def create_encoder(self, state_in, action_in, done_in, reuse):
         with tf.variable_scope("model"):
-            concat_input = tf.concat([state_in, action_in], axis=1)
+            concat_input = tf.concat([state_in, action_in, done_in], axis=1)
 
             hidden_1 = tf.layers.dense(
                 concat_input, self.h_size, activation=tf.nn.elu,
@@ -45,9 +49,9 @@ def create_encoder(self, state_in, action_in, reuse):
 
     def create_network(self):
         self.expert_estimate = self.create_encoder(
-            self.obs_in_expert, self.expert_action, False)
+            self.obs_in_expert, self.expert_action, self.done_expert, False)
         self.policy_estimate = self.create_encoder(
-            self.policy_model.vector_in, self.policy_model.selected_actions, True)
+            self.policy_model.vector_in, self.policy_model.selected_actions, self.done_policy, True)
         self.discriminator_score = tf.reshape(self.policy_estimate, [-1], name="GAIL_reward")
         self.intrinsic_reward = -tf.log(1.0 - self.discriminator_score + 1e-7)
 

ml-agents/mlagents/trainers/ppo/reward_signals/gail/signal.py

@@ -21,6 +21,7 @@ def evaluate(self, current_info, next_info):
         feed_dict = {self.policy.model.batch_size: len(next_info.vector_observations),
                      self.policy.model.sequence_length: 1}
         feed_dict = self.policy.fill_eval_dict(feed_dict, brain_info=current_info)
+        feed_dict[self.model.done_policy] = np.reshape(next_info.local_done, [-1, 1])
         if self.policy.use_continuous_act:
             feed_dict[self.policy.model.selected_actions] = next_info.previous_vector_actions
         else:
@@ -57,6 +58,9 @@ def update(self, policy_buffer, n_sequences, max_batches):
 
     def _update_batch(self, mini_batch_demo, mini_batch_policy):
         feed_dict = {}
+        feed_dict[self.model.done_expert] = mini_batch_demo['done'].reshape([-1, 1])
+        feed_dict[self.model.done_policy] = mini_batch_policy['done'].reshape([-1, 1])
+
         if self.policy.use_continuous_act:
             feed_dict[self.policy.model.selected_actions] = mini_batch_policy['actions'].reshape(
                 [-1, self.policy.model.act_size[0]])

ml-agents/mlagents/trainers/ppo/trainer.py

@@ -308,6 +308,7 @@ def add_experiences(
                     self.training_buffer[agent_id]['prev_action'].append(
                         stored_info.previous_vector_actions[idx])
                     self.training_buffer[agent_id]['masks'].append(1.0)
+                    self.training_buffer[agent_id]['done'].append(next_info.local_done[idx])
 
                     agent_rewards = None
                     for (scaled_reward, reward) in tmp_rewards_list: