Add Multi-GPU implementation for PPO (#2288)

Add MultiGpuPPOPolicy class and command line options to run multi-GPU training

Author: dongruoping

ml-agents/mlagents/trainers/learn.py

@@ -54,6 +54,7 @@ def run_training(
     lesson = int(run_options["--lesson"])
     fast_simulation = not bool(run_options["--slow"])
     no_graphics = run_options["--no-graphics"]
+    multi_gpu = run_options["--multi-gpu"]
     trainer_config_path = run_options["<trainer-config-path>"]
     sampler_file_path = (
         run_options["--sampler"] if run_options["--sampler"] != "None" else None
@@ -107,6 +108,7 @@ def run_training(
         lesson,
         run_seed,
         fast_simulation,
+        multi_gpu,
         sampler_manager,
         resampling_interval,
     )
@@ -292,6 +294,7 @@ def main():
       --docker-target-name=<dt>   Docker volume to store training-specific files [default: None].
       --no-graphics               Whether to run the environment in no-graphics mode [default: False].
       --debug                     Whether to run ML-Agents in debug mode with detailed logging [default: False].
+      --multi-gpu                Whether to use multiple GPU training [default: False].
     """
 
     options = docopt(_USAGE)

ml-agents/mlagents/trainers/models.py

@@ -601,6 +601,7 @@ def create_cc_actor_critic(
             hidden_policy,
             self.act_size[0],
             activation=None,
+            name="mu",
             kernel_initializer=c_layers.variance_scaling_initializer(factor=0.01),
         )
 
@@ -684,13 +685,14 @@ def create_dc_actor_critic(
             self.memory_out = tf.identity(memory_out, name="recurrent_out")
 
         policy_branches = []
-        for size in self.act_size:
+        for i, size in enumerate(self.act_size):
             policy_branches.append(
                 tf.layers.dense(
                     hidden,
                     size,
                     activation=None,
                     use_bias=False,
+                    name="policy_branch_" + str(i),
                     kernel_initializer=c_layers.variance_scaling_initializer(
                         factor=0.01
                     ),

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

@@ -138,5 +138,6 @@ def create_losses(
         )
 
     def create_ppo_optimizer(self):
-        optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
-        self.update_batch = optimizer.minimize(self.loss)
+        self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
+        self.grads = self.optimizer.compute_gradients(self.loss)
+        self.update_batch = self.optimizer.minimize(self.loss)

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

@@ -0,0 +1,140 @@
+import logging
+import numpy as np
+
+import tensorflow as tf
+from tensorflow.python.client import device_lib
+from mlagents.envs.timers import timed
+from mlagents.trainers.models import EncoderType
+from mlagents.trainers.ppo.policy import PPOPolicy
+from mlagents.trainers.ppo.models import PPOModel
+from mlagents.trainers.components.reward_signals.reward_signal_factory import (
+    create_reward_signal,
+)
+from mlagents.trainers.components.bc.module import BCModule
+
+# Variable scope in which created variables will be placed under
+TOWER_SCOPE_NAME = "tower"
+
+logger = logging.getLogger("mlagents.trainers")
+
+
+class MultiGpuPPOPolicy(PPOPolicy):
+    def __init__(self, seed, brain, trainer_params, is_training, load):
+        """
+        Policy for Proximal Policy Optimization Networks with multi-GPU training
+        :param seed: Random seed.
+        :param brain: Assigned Brain object.
+        :param trainer_params: Defined training parameters.
+        :param is_training: Whether the model should be trained.
+        :param load: Whether a pre-trained model will be loaded or a new one created.
+        """
+        super().__init__(seed, brain, trainer_params, is_training, load)
+
+        with self.graph.as_default():
+            avg_grads = self.average_gradients([t.grads for t in self.towers])
+            self.update_batch = self.model.optimizer.apply_gradients(avg_grads)
+
+        self.update_dict = {"update_batch": self.update_batch}
+        self.update_dict.update(
+            {
+                "value_loss_" + str(i): self.towers[i].value_loss
+                for i in range(len(self.towers))
+            }
+        )
+        self.update_dict.update(
+            {
+                "policy_loss_" + str(i): self.towers[i].policy_loss
+                for i in range(len(self.towers))
+            }
+        )
+
+    def create_model(self, brain, trainer_params, reward_signal_configs, seed):
+        """
+        Create PPO models, one on each device
+        :param brain: Assigned Brain object.
+        :param trainer_params: Defined training parameters.
+        :param reward_signal_configs: Reward signal config
+        :param seed: Random seed.
+        """
+        self.devices = get_devices()
+        self.towers = []
+        with self.graph.as_default():
+            with tf.variable_scope(TOWER_SCOPE_NAME, reuse=tf.AUTO_REUSE):
+                for device in self.devices:
+                    with tf.device(device):
+                        self.towers.append(
+                            PPOModel(
+                                brain=brain,
+                                lr=float(trainer_params["learning_rate"]),
+                                h_size=int(trainer_params["hidden_units"]),
+                                epsilon=float(trainer_params["epsilon"]),
+                                beta=float(trainer_params["beta"]),
+                                max_step=float(trainer_params["max_steps"]),
+                                normalize=trainer_params["normalize"],
+                                use_recurrent=trainer_params["use_recurrent"],
+                                num_layers=int(trainer_params["num_layers"]),
+                                m_size=self.m_size,
+                                seed=seed,
+                                stream_names=list(reward_signal_configs.keys()),
+                                vis_encode_type=EncoderType(
+                                    trainer_params.get("vis_encode_type", "simple")
+                                ),
+                            )
+                        )
+                        self.towers[-1].create_ppo_optimizer()
+            self.model = self.towers[0]
+
+    @timed
+    def update(self, mini_batch, num_sequences):
+        """
+        Updates model using buffer.
+        :param n_sequences: Number of trajectories in batch.
+        :param mini_batch: Experience batch.
+        :return: Output from update process.
+        """
+        feed_dict = {}
+
+        device_batch_size = num_sequences // len(self.devices)
+        device_batches = []
+        for i in range(len(self.devices)):
+            device_batches.append(
+                {k: v[i : i + device_batch_size] for (k, v) in mini_batch.items()}
+            )
+
+        for batch, tower in zip(device_batches, self.towers):
+            feed_dict.update(self.construct_feed_dict(tower, batch, num_sequences))
+
+        out = self._execute_model(feed_dict, self.update_dict)
+        run_out = {}
+        run_out["value_loss"] = np.mean(
+            [out["value_loss_" + str(i)] for i in range(len(self.towers))]
+        )
+        run_out["policy_loss"] = np.mean(
+            [out["policy_loss_" + str(i)] for i in range(len(self.towers))]
+        )
+        run_out["update_batch"] = out["update_batch"]
+        return run_out
+
+    def average_gradients(self, tower_grads):
+        """
+        Average gradients from all towers
+        :param tower_grads: Gradients from all towers
+        """
+        average_grads = []
+        for grad_and_vars in zip(*tower_grads):
+            grads = [g for g, _ in grad_and_vars if g is not None]
+            if not grads:
+                continue
+            avg_grad = tf.reduce_mean(tf.stack(grads), 0)
+            var = grad_and_vars[0][1]
+            average_grads.append((avg_grad, var))
+        return average_grads
+
+
+def get_devices():
+    """
+    Get all available GPU devices
+    """
+    local_device_protos = device_lib.list_local_devices()
+    devices = [x.name for x in local_device_protos if x.device_type == "GPU"]
+    return devices

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

@@ -30,27 +30,10 @@ def __init__(self, seed, brain, trainer_params, is_training, load):
 
         reward_signal_configs = trainer_params["reward_signals"]
 
+        self.create_model(brain, trainer_params, reward_signal_configs, seed)
+
         self.reward_signals = {}
         with self.graph.as_default():
-            self.model = PPOModel(
-                brain,
-                lr=float(trainer_params["learning_rate"]),
-                h_size=int(trainer_params["hidden_units"]),
-                epsilon=float(trainer_params["epsilon"]),
-                beta=float(trainer_params["beta"]),
-                max_step=float(trainer_params["max_steps"]),
-                normalize=trainer_params["normalize"],
-                use_recurrent=trainer_params["use_recurrent"],
-                num_layers=int(trainer_params["num_layers"]),
-                m_size=self.m_size,
-                seed=seed,
-                stream_names=list(reward_signal_configs.keys()),
-                vis_encode_type=EncoderType(
-                    trainer_params.get("vis_encode_type", "simple")
-                ),
-            )
-            self.model.create_ppo_optimizer()
-
             # Create reward signals
             for reward_signal, config in reward_signal_configs.items():
                 self.reward_signals[reward_signal] = create_reward_signal(
@@ -102,6 +85,34 @@ def __init__(self, seed, brain, trainer_params, is_training, load):
             "update_batch": self.model.update_batch,
         }
 
+    def create_model(self, brain, trainer_params, reward_signal_configs, seed):
+        """
+        Create PPO model
+        :param brain: Assigned Brain object.
+        :param trainer_params: Defined training parameters.
+        :param reward_signal_configs: Reward signal config
+        :param seed: Random seed.
+        """
+        with self.graph.as_default():
+            self.model = PPOModel(
+                brain=brain,
+                lr=float(trainer_params["learning_rate"]),
+                h_size=int(trainer_params["hidden_units"]),
+                epsilon=float(trainer_params["epsilon"]),
+                beta=float(trainer_params["beta"]),
+                max_step=float(trainer_params["max_steps"]),
+                normalize=trainer_params["normalize"],
+                use_recurrent=trainer_params["use_recurrent"],
+                num_layers=int(trainer_params["num_layers"]),
+                m_size=self.m_size,
+                seed=seed,
+                stream_names=list(reward_signal_configs.keys()),
+                vis_encode_type=EncoderType(
+                    trainer_params.get("vis_encode_type", "simple")
+                ),
+            )
+            self.model.create_ppo_optimizer()
+
     @timed
     def evaluate(self, brain_info):
         """
@@ -143,58 +154,62 @@ def update(self, mini_batch, num_sequences):
         :param mini_batch: Experience batch.
         :return: Output from update process.
         """
+        feed_dict = self.construct_feed_dict(self.model, mini_batch, num_sequences)
+        run_out = self._execute_model(feed_dict, self.update_dict)
+        return run_out
+
+    def construct_feed_dict(self, model, mini_batch, num_sequences):
         feed_dict = {
-            self.model.batch_size: num_sequences,
-            self.model.sequence_length: self.sequence_length,
-            self.model.mask_input: mini_batch["masks"].flatten(),
-            self.model.advantage: mini_batch["advantages"].reshape([-1, 1]),
-            self.model.all_old_log_probs: mini_batch["action_probs"].reshape(
-                [-1, sum(self.model.act_size)]
+            model.batch_size: num_sequences,
+            model.sequence_length: self.sequence_length,
+            model.mask_input: mini_batch["masks"].flatten(),
+            model.advantage: mini_batch["advantages"].reshape([-1, 1]),
+            model.all_old_log_probs: mini_batch["action_probs"].reshape(
+                [-1, sum(model.act_size)]
             ),
         }
         for name in self.reward_signals:
-            feed_dict[self.model.returns_holders[name]] = mini_batch[
+            feed_dict[model.returns_holders[name]] = mini_batch[
                 "{}_returns".format(name)
             ].flatten()
-            feed_dict[self.model.old_values[name]] = mini_batch[
+            feed_dict[model.old_values[name]] = mini_batch[
                 "{}_value_estimates".format(name)
             ].flatten()
 
         if self.use_continuous_act:
-            feed_dict[self.model.output_pre] = mini_batch["actions_pre"].reshape(
-                [-1, self.model.act_size[0]]
+            feed_dict[model.output_pre] = mini_batch["actions_pre"].reshape(
+                [-1, model.act_size[0]]
             )
-            feed_dict[self.model.epsilon] = mini_batch["random_normal_epsilon"].reshape(
-                [-1, self.model.act_size[0]]
+            feed_dict[model.epsilon] = mini_batch["random_normal_epsilon"].reshape(
+                [-1, model.act_size[0]]
             )
         else:
-            feed_dict[self.model.action_holder] = mini_batch["actions"].reshape(
-                [-1, len(self.model.act_size)]
+            feed_dict[model.action_holder] = mini_batch["actions"].reshape(
+                [-1, len(model.act_size)]
             )
             if self.use_recurrent:
-                feed_dict[self.model.prev_action] = mini_batch["prev_action"].reshape(
-                    [-1, len(self.model.act_size)]
+                feed_dict[model.prev_action] = mini_batch["prev_action"].reshape(
+                    [-1, len(model.act_size)]
                 )
-            feed_dict[self.model.action_masks] = mini_batch["action_mask"].reshape(
+            feed_dict[model.action_masks] = mini_batch["action_mask"].reshape(
                 [-1, sum(self.brain.vector_action_space_size)]
             )
         if self.use_vec_obs:
-            feed_dict[self.model.vector_in] = mini_batch["vector_obs"].reshape(
+            feed_dict[model.vector_in] = mini_batch["vector_obs"].reshape(
                 [-1, self.vec_obs_size]
             )
-        if self.model.vis_obs_size > 0:
-            for i, _ in enumerate(self.model.visual_in):
+        if model.vis_obs_size > 0:
+            for i, _ in enumerate(model.visual_in):
                 _obs = mini_batch["visual_obs%d" % i]
                 if self.sequence_length > 1 and self.use_recurrent:
                     (_batch, _seq, _w, _h, _c) = _obs.shape
-                    feed_dict[self.model.visual_in[i]] = _obs.reshape([-1, _w, _h, _c])
+                    feed_dict[model.visual_in[i]] = _obs.reshape([-1, _w, _h, _c])
                 else:
-                    feed_dict[self.model.visual_in[i]] = _obs
+                    feed_dict[model.visual_in[i]] = _obs
         if self.use_recurrent:
             mem_in = mini_batch["memory"][:, 0, :]
-            feed_dict[self.model.memory_in] = mem_in
-        run_out = self._execute_model(feed_dict, self.update_dict)
-        return run_out
+            feed_dict[model.memory_in] = mem_in
+        return feed_dict
 
     def get_value_estimates(
         self, brain_info: BrainInfo, idx: int, done: bool

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

@@ -11,6 +11,7 @@
 from mlagents.envs import AllBrainInfo, BrainInfo
 from mlagents.trainers.buffer import Buffer
 from mlagents.trainers.ppo.policy import PPOPolicy
+from mlagents.trainers.ppo.multi_gpu_policy import MultiGpuPPOPolicy, get_devices
 from mlagents.trainers.trainer import Trainer, UnityTrainerException
 from mlagents.envs.action_info import ActionInfoOutputs
 
@@ -21,7 +22,15 @@ class PPOTrainer(Trainer):
     """The PPOTrainer is an implementation of the PPO algorithm."""
 
     def __init__(
-        self, brain, reward_buff_cap, trainer_parameters, training, load, seed, run_id
+        self,
+        brain,
+        reward_buff_cap,
+        trainer_parameters,
+        training,
+        load,
+        seed,
+        run_id,
+        multi_gpu,
     ):
         """
         Responsible for collecting experiences and training PPO model.
@@ -65,7 +74,14 @@ def __init__(
             )
 
         self.step = 0
-        self.policy = PPOPolicy(seed, brain, trainer_parameters, self.is_training, load)
+        if multi_gpu and len(get_devices()) > 1:
+            self.policy = MultiGpuPPOPolicy(
+                seed, brain, trainer_parameters, self.is_training, load
+            )
+        else:
+            self.policy = PPOPolicy(
+                seed, brain, trainer_parameters, self.is_training, load
+            )
 
         stats = defaultdict(list)
         # collected_rewards is a dictionary from name of reward signal to a dictionary of agent_id to cumulative reward