第一版稳定FusionPPO更新完毕，修复底层bug

1. 现在可以放心在有限制的情况下使用FusionPPO。该算法在简单的POMDP提升效果并不明显甚至会变慢，但是在复杂环境中，可能只有这个算法有效。
2. 修复底层并行bug，现在你可以随心所欲的并行，只需要花一点点时间。

Author: yangtao121

AquaML/rlalgo/BaseRLAlgo.py

@@ -724,8 +724,8 @@ def create_gaussian_exploration_policy(self):
             self.log_std = DataUnit(name=self.name + '_log_std', dtype=np.float32,
                                     shape=self.rl_io_info.actor_out_info['action'],
                                     level=self.level, computer_type=self._computer_type)
-
-            self.log_std.set_value(np.zeros(self.rl_io_info.actor_out_info['action'], dtype=np.float32) - 0.5)
+            if self.level == 0:
+                self.log_std.set_value(np.zeros(self.rl_io_info.actor_out_info['action'], dtype=np.float32) - 0.5)
             self.tf_log_std = tf.Variable(self.log_std.buffer, trainable=True)
             self._explore_dict = {'log_std': self.tf_log_std}
 

AquaML/rlalgo/FusionPPO.py

@@ -83,6 +83,7 @@ def __init__(self,
                     fusion_flag = True
                     break
                 idx += 1
+            # self.fusion_value_idx += 1
             if not fusion_flag:
                 raise ValueError('Fusion value must be in actor output. '
                                  'Please check your actor output.')
@@ -99,6 +100,8 @@ def __init__(self,
         # initialize actor
         # self.initialize_model_weights(self.actor)
 
+        self._sync_model_dict['actor'] = self.actor
+
         # create optimizer
         if self.level == 0:
             self.create_optimizer(name='actor', optimizer=self.actor.optimizer, lr=self.actor.learning_rate)
@@ -161,6 +164,42 @@ def train_actor(self,
         actor_grad = tape.gradient(loss, self.actor.trainable_variables)
         self.actor_optimizer.apply_gradients(zip(actor_grad, self.actor.trainable_variables))
 
+        # with tf.GradientTape() as tape:
+        #     out = self.resample_log_prob(actor_obs, action)
+        #     log_prob = out[0]
+        #     fusion_value = out[self.fusion_value_idx]
+        #
+        #     ratio = tf.exp(log_prob - old_log_prob)
+        #
+        #     actor_surrogate = tf.minimum(
+        #         ratio * advantage,
+        #         tf.clip_by_value(ratio, 1 - epsilon, 1 + epsilon) * advantage,
+        #     )
+        #
+        #     entropy = -log_prob
+        #     fusion_value_d = tf.square(fusion_value - target)
+        #
+        #     normalized_surrogate_loss = tf.reduce_mean(tf.math.l2_normalize(actor_surrogate, axis=0))
+        #
+        #     normalized_entropy_loss = tf.reduce_mean(tf.math.l2_normalize(entropy, axis=0))
+        #
+        #     normalized_fusion_value_loss = tf.reduce_mean(tf.math.l2_normalize(fusion_value_d, axis=0))
+        #
+        #     normalized_loss = -normalized_surrogate_loss + lam * normalized_fusion_value_loss - entropy_coefficient * normalized_entropy_loss
+        #
+        # normalized_actor_grad = tape.gradient(normalized_loss, self.actor.trainable_variables)
+        # self.actor_optimizer.apply_gradients(zip(normalized_actor_grad, self.actor.trainable_variables))
+
+        # dic = {
+        #     'actor_surrogate_loss': tf.reduce_mean(actor_surrogate),
+        #     'actor_loss': normalized_loss,
+        #     'fusion_value_loss': tf.reduce_mean(fusion_value_d),
+        #     'entropy_loss': tf.reduce_mean(entropy),
+        #     # 'normalized_actor_loss': normalized_loss,
+        #     'normalized_actor_surrogate_loss': normalized_surrogate_loss,
+        #     'normalized_fusion_value_loss': normalized_fusion_value_loss,
+        #     'normalized_entropy_loss': normalized_entropy_loss,
+        # }
         dic = {
             'actor_surrogate_loss': actor_surrogate_loss,
             'actor_loss': loss,
@@ -252,72 +291,142 @@ def _optimize_(self):
             else:
                 for idx in self.expand_dims_idx:
                     actor_obs[idx] = tf.expand_dims(actor_obs[idx], axis=1)
+        info_list = []
+        buffer_size = train_actor_input['actor_obs'][0].shape[0]
+
+        if self.hyper_parameters.batch_trajectory:
+            critic_batch_steps = self.hyper_parameters.batch_size * train_actor_input['actor_obs'][0].shape[1]
+        else:
+            critic_batch_steps = self.hyper_parameters.batch_size
+
+        critic_buffer_size = self.hyper_parameters.buffer_size
 
         for _ in range(self.hyper_parameters.update_times):
             # fusion ppo firstly update critic
-            for _ in range(self.hyper_parameters.update_critic_times):
-                start_index = 0
-                end_index = 0
+            start_index = 0
+            end_index = 0
+            critic_start_index = 0
+            while end_index < buffer_size:
+                end_index = min(start_index + self.hyper_parameters.batch_size,
+                                buffer_size)
+                critic_end_index = min(critic_start_index + critic_batch_steps, critic_buffer_size)
                 critic_optimize_info_list = []
+                actor_optimize_info_list = []
+                batch_train_actor_input = self.get_batch_data(train_actor_input, start_index, end_index)
+                batch_train_critic_input = self.get_batch_data(train_critic_input, critic_start_index, critic_end_index)
+                start_index = end_index
+                critic_start_index = critic_end_index
                 for _ in range(self.hyper_parameters.update_critic_times):
-                    while end_index < self.hyper_parameters.buffer_size:
-                        end_index = min(start_index + self.hyper_parameters.batch_size,
-                                        self.hyper_parameters.buffer_size)
+                    critic_optimize_info = self.train_critic(
+                        critic_obs=batch_train_critic_input['critic_obs'],
+                        target=batch_train_critic_input['target'],
+                    )
+                    critic_optimize_info_list.append(critic_optimize_info)
 
-                        batch_train_critic_input = self.get_batch_data(train_critic_input, start_index, end_index)
+                critic_value = self.critic(*batch_train_critic_input['critic_obs'])
+                critic_value_target = tf.math.reduce_mean(tf.square(critic_value - batch_train_critic_input['target']))
 
-                        start_index = end_index
+                out = self.resample_log_prob(batch_train_actor_input['actor_obs'], batch_train_actor_input['action'])
+                fusion_value = out[self.fusion_value_idx]
 
-                        critic_optimize_info = self.train_critic(
-                            critic_obs=batch_train_critic_input['critic_obs'],
-                            target=batch_train_critic_input['target'],
-                        )
-                        critic_optimize_info_list.append(critic_optimize_info)
-                    critic_optimize_info = self.cal_average_batch_dict(critic_optimize_info_list)
+                # fusion_value = tf.reshape(fusion_value, critic_value.shape)
+                critic_value = tf.reshape(critic_value, shape=fusion_value.shape)
 
-            # fusion ppo secondly update actor
-            # compute lam
-            critic_value = self.critic(*critic_obs)
-            critic_value_target = tf.reduce_mean(tf.square(critic_value - target))
+                fusion_value_critic = tf.math.reduce_mean(tf.square(fusion_value - critic_value))
 
-            out = self.resample_log_prob(actor_obs, train_actor_input['action'])
-            fusion_value = out[self.fusion_value_idx]
-
-            fusion_value = tf.reshape(fusion_value, critic_value.shape)
-
-            fusion_value_critic = tf.reduce_mean(tf.square(fusion_value - critic_value))
+                # distance = tf.sqrt(critic_value_target) + tf.sqrt(fusion_value_critic)
+                distance = critic_value_target + fusion_value_critic
 
-            distance = tf.sqrt(critic_value_target) + tf.sqrt(fusion_value_critic)
+                lam = 1. / distance
+                lam = tf.clip_by_value(lam, 0, 0.2)
+                # lam = 1
+                lam = 0
 
-            batch_size = train_actor_input['actor_obs'][0].shape[0]
-
-            lam = 1. / distance
-            # lam = 0
-            for _ in range(self.hyper_parameters.update_actor_times):
-                start_index = 0
-                end_index = 0
-                actor_optimize_info_list = []
                 for _ in range(self.hyper_parameters.update_actor_times):
-                    while end_index < batch_size:
-                        end_index = min(start_index + self.hyper_parameters.batch_size,
-                                        batch_size)
-
-                        batch_train_actor_input = self.get_batch_data(train_actor_input, start_index, end_index)
-
-                        start_index = end_index
-
-                        actor_optimize_info = self.train_actor(
-                            actor_obs=batch_train_actor_input['actor_obs'],
-                            advantage=batch_train_actor_input['advantage'],
-                            old_log_prob=batch_train_actor_input['old_log_prob'],
-                            action=batch_train_actor_input['action'],
-                            target=batch_train_actor_input['target'],
-                            lam=lam,
-                            epsilon=tf.cast(self.hyper_parameters.epsilon, dtype=tf.float32),
-                            entropy_coefficient=tf.cast(self.hyper_parameters.entropy_coeff, dtype=tf.float32),
-                        )
-                        actor_optimize_info_list.append(actor_optimize_info)
-                    actor_optimize_info = self.cal_average_batch_dict(actor_optimize_info_list)
-
-        return_dict = {**critic_optimize_info, **actor_optimize_info, 'lam': lam}
-        return return_dict
+                    actor_optimize_info = self.train_actor(
+                        actor_obs=batch_train_actor_input['actor_obs'],
+                        advantage=batch_train_actor_input['advantage'],
+                        old_log_prob=batch_train_actor_input['old_log_prob'],
+                        action=batch_train_actor_input['action'],
+                        target=batch_train_actor_input['target'],
+                        lam=lam,
+                        epsilon=tf.cast(self.hyper_parameters.epsilon, dtype=tf.float32),
+                        entropy_coefficient=tf.cast(self.hyper_parameters.entropy_coeff, dtype=tf.float32),
+                    )
+                    actor_optimize_info_list.append(actor_optimize_info)
+                critic_optimize_info = self.cal_average_batch_dict(critic_optimize_info_list)
+                actor_optimize_info = self.cal_average_batch_dict(actor_optimize_info_list)
+                info = {**critic_optimize_info, **actor_optimize_info, 'lam': lam}
+                info_list.append(info)
+
+            info = self.cal_average_batch_dict(info_list)
+
+            return info
+
+        #     for _ in range(self.hyper_parameters.update_critic_times):
+        #         start_index = 0
+        #         end_index = 0
+        #         critic_optimize_info_list = []
+        #         for _ in range(self.hyper_parameters.update_critic_times):
+        #             while end_index < self.hyper_parameters.buffer_size:
+        #                 end_index = min(start_index + self.hyper_parameters.batch_size,
+        #                                 self.hyper_parameters.buffer_size)
+        #
+        #                 batch_train_critic_input = self.get_batch_data(train_critic_input, start_index, end_index)
+        #
+        #                 start_index = end_index
+        #
+        #                 critic_optimize_info = self.train_critic(
+        #                     critic_obs=batch_train_critic_input['critic_obs'],
+        #                     target=batch_train_critic_input['target'],
+        #                 )
+        #                 critic_optimize_info_list.append(critic_optimize_info)
+        #
+        #
+        #
+        #     # fusion ppo secondly update actor
+        #     # compute lam
+        #     critic_value = self.critic(*critic_obs)
+        #     critic_value_target = tf.reduce_mean(tf.square(critic_value - target))
+        #
+        #     out = self.resample_log_prob(actor_obs, train_actor_input['action'])
+        #     fusion_value = out[self.fusion_value_idx]
+        #
+        #     fusion_value = tf.reshape(fusion_value, critic_value.shape)
+        #
+        #     fusion_value_critic = tf.reduce_mean(tf.square(fusion_value - critic_value))
+        #
+        #     distance = tf.sqrt(critic_value_target) + tf.sqrt(fusion_value_critic)
+        #
+        #     batch_size = train_actor_input['actor_obs'][0].shape[0]
+        #
+        #     lam = 1. / distance
+        #     # lam = 0
+        #     for _ in range(self.hyper_parameters.update_actor_times):
+        #         start_index = 0
+        #         end_index = 0
+        #         actor_optimize_info_list = []
+        #         for _ in range(self.hyper_parameters.update_actor_times):
+        #             while end_index < batch_size:
+        #                 end_index = min(start_index + self.hyper_parameters.batch_size,
+        #                                 batch_size)
+        #
+        #                 batch_train_actor_input = self.get_batch_data(train_actor_input, start_index, end_index)
+        #
+        #                 start_index = end_index
+        #
+        #                 actor_optimize_info = self.train_actor(
+        #                     actor_obs=batch_train_actor_input['actor_obs'],
+        #                     advantage=batch_train_actor_input['advantage'],
+        #                     old_log_prob=batch_train_actor_input['old_log_prob'],
+        #                     action=batch_train_actor_input['action'],
+        #                     target=batch_train_actor_input['target'],
+        #                     lam=lam,
+        #                     epsilon=tf.cast(self.hyper_parameters.epsilon, dtype=tf.float32),
+        #                     entropy_coefficient=tf.cast(self.hyper_parameters.entropy_coeff, dtype=tf.float32),
+        #                 )
+        #                 actor_optimize_info_list.append(actor_optimize_info)
+        #             actor_optimize_info = self.cal_average_batch_dict(actor_optimize_info_list)
+        #
+        # return_dict = {**critic_optimize_info, **actor_optimize_info, 'lam': lam}
+        # return return_dict

AquaML/starter/RLTaskStarter.py

@@ -137,13 +137,13 @@ def _run_(self):
 
     def _run_mpi_(self):
         for i in range(self.max_epochs):
-            if self.thread_id == 0:
+            if self.level == 0:
                 self.algo.sync()
             else:
                 pass
             self.mpi_comm.Barrier()
 
-            if self.thread_id > 0:
+            if self.level == 1:
                 self.algo.sync()
                 self.algo.worker.roll(self.roll_out_length)
                 self.roll_out_length = self.update_interval

Tutorial/FPPO2/cache/actor.h5

@@ -3,6 +3,14 @@
 
 The environment we use is the POMDP Pendulum-v1.
 """
+import sys
+sys.path.append('..')
+from AquaML.Tool import allocate_gpu
+from mpi4py import MPI
+
+# get group communicator
+comm = MPI.COMM_WORLD
+allocate_gpu(comm)
 
 import tensorflow as tf
 from AquaML.DataType import DataInfo
@@ -24,12 +32,14 @@ def __init__(self):
         self.dense2 = tf.keras.layers.Dense(64, activation='relu')
 
         self.action_dense = tf.keras.layers.Dense(64, activation='relu')
+        self.action_dense2 = tf.keras.layers.Dense(64, activation='relu')
         self.action_layer = tf.keras.layers.Dense(1, activation='tanh')
 
         self.value_dense = tf.keras.layers.Dense(64, activation='relu')
+        self.value_dense2 = tf.keras.layers.Dense(64, activation='relu')
         self.value_layer = tf.keras.layers.Dense(1)
 
-        self.learning_rate = 2e-4
+        self.learning_rate = 2e-3
 
         self.rnn_flag = True
 
@@ -45,8 +55,10 @@ def call(self, vel, hidden1, hidden2):
         whole_seq, last_seq, hidden_state = self.lstm(vel, hidden_states)
         x = self.dense2(whole_seq)
         action_x = self.action_dense(x)
+        action_x = self.action_dense2(action_x)
         action = self.action_layer(action_x)
         value_x = self.value_dense(x)
+        value_x = self.value_dense2(value_x)
         value = self.value_layer(value_x)
 
         return (action, value, last_seq, hidden_state)
@@ -136,11 +148,11 @@ def close(self):
 fusion_ppo_parameter = FusionPPO_parameter(
     epoch_length=200,
     n_epochs=2000,
-    total_steps=4000,
-    batch_size=256,
+    total_steps=6000,
+    batch_size=20,
     update_times=4,
-    update_actor_times=2,
-    update_critic_times=4,
+    update_actor_times=6,
+    update_critic_times=6,
     gamma=0.99,
     epsilon=0.2,
     lambada=0.95,
@@ -157,7 +169,8 @@ def close(self):
     model_class_dict=model_class_dict,
     algo=FusionPPO,
     algo_hyperparameter=fusion_ppo_parameter,
-    name='FPPO'
+    name='FPPO2',
+    mpi_comm=comm,
 )
 
 starter.run()