GAIL no longer uses placeholders from Policy

Author: Ervin Teng

ml-agents/mlagents/trainers/components/reward_signals/curiosity/signal.py

@@ -57,7 +57,6 @@ def evaluate(
             return []
 
         feed_dict = {}
-
         for i, _ in enumerate(current_info.visual_observations):
             feed_dict[self.model.visual_in[i]] = current_info.visual_observations[i]
         if self.policy.use_vec_obs:

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

@@ -63,15 +63,29 @@ def make_inputs(self) -> None:
 
         if self.policy_model.brain.vector_action_space_type == "continuous":
             action_length = self.policy_model.act_size[0]
+            self.action_in_policy = tf.placeholder(
+                shape=[None, action_length], dtype=tf.float32
+            )
             self.action_in_expert = tf.placeholder(
                 shape=[None, action_length], dtype=tf.float32
             )
             self.expert_action = tf.identity(self.action_in_expert)
+            self.policy_action = tf.identity(self.action_in_policy)
         else:
             action_length = len(self.policy_model.act_size)
+            self.action_in_policy = tf.placeholder(
+                shape=[None, action_length], dtype=tf.int32
+            )
             self.action_in_expert = tf.placeholder(
                 shape=[None, action_length], dtype=tf.int32
             )
+            self.policy_action = tf.concat(
+                [
+                    tf.one_hot(self.action_in_policy[:, i], act_size)
+                    for i, act_size in enumerate(self.policy_model.act_size)
+                ],
+                axis=1,
+            )
             self.expert_action = tf.concat(
                 [
                     tf.one_hot(self.action_in_expert[:, i], act_size)
@@ -84,6 +98,9 @@ def make_inputs(self) -> None:
         encoded_expert_list = []
 
         if self.policy_model.vec_obs_size > 0:
+            self.vector_in = tf.placeholder(
+                shape=[None, self.policy_model.vec_obs_size], dtype=tf.float32
+            )
             self.obs_in_expert = tf.placeholder(
                 shape=[None, self.policy_model.vec_obs_size], dtype=tf.float32
             )
@@ -92,26 +109,33 @@ def make_inputs(self) -> None:
                     self.policy_model.normalize_vector_obs(self.obs_in_expert)
                 )
                 encoded_policy_list.append(
-                    self.policy_model.normalize_vector_obs(self.policy_model.vector_in)
+                    self.policy_model.normalize_vector_obs(self.vector_in)
                 )
             else:
                 encoded_expert_list.append(self.obs_in_expert)
-                encoded_policy_list.append(self.policy_model.vector_in)
+                encoded_policy_list.append(self.vector_in)
 
         if self.policy_model.vis_obs_size > 0:
             self.expert_visual_in: List[tf.Tensor] = []
+            self.visual_in: List[tf.Tensor] = []
             visual_policy_encoders = []
             visual_expert_encoders = []
             for i in range(self.policy_model.vis_obs_size):
-                # Create input ops for next (t+1) visual observations.
+                # Create input ops for visual observations.
                 visual_input = self.policy_model.create_visual_input(
                     self.policy_model.brain.camera_resolutions[i],
-                    name="visual_observation_" + str(i),
+                    name="gail_visual_observation_" + str(i),
+                )
+                self.visual_in.append(visual_input)
+                # Create input ops for next (t+1) visual observations.
+                ex_visual_input = self.policy_model.create_visual_input(
+                    self.policy_model.brain.camera_resolutions[i],
+                    name="expert_visual_observation_" + str(i),
                 )
-                self.expert_visual_in.append(visual_input)
+                self.expert_visual_in.append(ex_visual_input)
 
                 encoded_policy_visual = self.policy_model.create_visual_observation_encoder(
-                    self.policy_model.visual_in[i],
+                    self.visual_in[i],
                     self.encoding_size,
                     LearningModel.swish,
                     1,
@@ -217,10 +241,7 @@ def create_network(self) -> None:
             self.encoded_expert, self.expert_action, self.done_expert, reuse=False
         )
         self.policy_estimate, self.z_mean_policy, _ = self.create_encoder(
-            self.encoded_policy,
-            self.policy_model.selected_actions,
-            self.done_policy,
-            reuse=True,
+            self.encoded_policy, self.policy_action, self.done_policy, reuse=True
         )
         self.discriminator_score = tf.reshape(
             self.policy_estimate, [-1], name="GAIL_reward"
@@ -233,11 +254,7 @@ def create_gradient_magnitude(self) -> tf.Tensor:
         for off-policy. Compute gradients w.r.t randomly interpolated input.
         """
         expert = [self.encoded_expert, self.expert_action, self.done_expert]
-        policy = [
-            self.encoded_policy,
-            self.policy_model.selected_actions,
-            self.done_policy,
-        ]
+        policy = [self.encoded_policy, self.policy_action, self.done_policy]
         interp = []
         for _expert_in, _policy_in in zip(expert, policy):
             alpha = tf.random_uniform(tf.shape(_expert_in))

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

@@ -67,24 +67,19 @@ def evaluate(
     ) -> RewardSignalResult:
         if len(current_info.agents) == 0:
             return []
-
-        feed_dict: Dict[tf.Tensor, Any] = {
-            self.policy.model.batch_size: len(next_info.vector_observations),
-            self.policy.model.sequence_length: 1,
-        }
+        feed_dict: Dict[tf.Tensor, Any] = {}
         if self.model.use_vail:
             feed_dict[self.model.use_noise] = [0]
+        for i, _ in enumerate(current_info.visual_observations):
+            feed_dict[self.model.visual_in[i]] = current_info.visual_observations[i]
+        if self.policy.use_vec_obs:
+            feed_dict[self.model.vector_in] = current_info.vector_observations
 
-        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
+            feed_dict[self.model.action_in_policy] = next_info.previous_vector_actions
         else:
-            feed_dict[
-                self.policy.model.action_holder
-            ] = next_info.previous_vector_actions
+            feed_dict[self.model.action_in_policy] = next_info.previous_vector_actions
         unscaled_reward = self.policy.sess.run(
             self.model.intrinsic_reward, feed_dict=feed_dict
         )
@@ -134,30 +129,30 @@ def prepare_update(
             feed_dict[self.model.use_noise] = [1]
 
         if self.policy.use_continuous_act:
-            feed_dict[self.policy.model.selected_actions] = mini_batch_policy[
+            feed_dict[self.model.action_in_policy] = mini_batch_policy[
                 "actions"
             ].reshape([-1, self.policy.model.act_size[0]])
             feed_dict[self.model.action_in_expert] = mini_batch_demo["actions"].reshape(
                 [-1, self.policy.model.act_size[0]]
             )
         else:
-            feed_dict[self.policy.model.action_holder] = mini_batch_policy[
+            feed_dict[self.model.action_in_policy] = mini_batch_policy[
                 "actions"
             ].reshape([-1, len(self.policy.model.act_size)])
             feed_dict[self.model.action_in_expert] = mini_batch_demo["actions"].reshape(
                 [-1, len(self.policy.model.act_size)]
             )
 
         if self.policy.use_vis_obs > 0:
-            for i in range(len(self.policy.model.visual_in)):
+            for i in range(len(self.model.visual_in)):
                 policy_obs = mini_batch_policy["visual_obs%d" % i]
                 if self.policy.sequence_length > 1 and self.policy.use_recurrent:
                     (_batch, _seq, _w, _h, _c) = policy_obs.shape
-                    feed_dict[self.policy.model.visual_in[i]] = policy_obs.reshape(
+                    feed_dict[self.model.visual_in[i]] = policy_obs.reshape(
                         [-1, _w, _h, _c]
                     )
                 else:
-                    feed_dict[self.policy.model.visual_in[i]] = policy_obs
+                    feed_dict[self.model.visual_in[i]] = policy_obs
 
                 demo_obs = mini_batch_demo["visual_obs%d" % i]
                 if self.policy.sequence_length > 1 and self.policy.use_recurrent:
@@ -168,9 +163,9 @@ def prepare_update(
                 else:
                     feed_dict[self.model.expert_visual_in[i]] = demo_obs
         if self.policy.use_vec_obs:
-            feed_dict[self.policy.model.vector_in] = mini_batch_policy[
-                "vector_obs"
-            ].reshape([-1, self.policy.vec_obs_size])
+            feed_dict[self.model.vector_in] = mini_batch_policy["vector_obs"].reshape(
+                [-1, self.policy.vec_obs_size]
+            )
             feed_dict[self.model.obs_in_expert] = mini_batch_demo["vector_obs"].reshape(
                 [-1, self.policy.vec_obs_size]
             )