moved par_opt out of evaluator

python/ray/rllib/ppo/ppo.py

@@ -16,6 +16,9 @@
 from ray.rllib.agent import Agent
 from ray.rllib.utils import FilterManager
 from ray.rllib.ppo.ppo_evaluator import PPOEvaluator
+from ray.rllib.optimizers.multi_gpu_impl import LocalSyncParallelOptimizer
+from ray.rllib.utils.common_policy_evaluator import CommonPolicyEvaluator, \
+    collect_metrics
 from ray.rllib.ppo.rollout import collect_samples
 
 
@@ -109,6 +112,22 @@ def _init(self):
         self.kl_coeff_val = self.config["kl_coeff"]
         self.local_evaluator = PPOEvaluator(
             self.registry, self.env_creator, self.config, self.logdir, False)
+
+        main_thread_scope = tf.get_variable_scope()
+
+        with tf.variable_scope(main_thread_scope, reuse=tf.AUTO_REUSE):
+            self.par_opt = LocalSyncParallelOptimizer(
+                tf.train.AdamOptimizer(self.config["sgd_stepsize"]),
+                self.local_evaluator.devices,
+                self.local_evaluator.inputs,
+                self.local_evaluator.per_device_batch_size,
+                self.local_evaluator.build_loss,
+                self.logdir)
+            self.local_evaluator.init_extra_ops(
+                self.par_opt.get_device_losses())
+
+        self.local_evaluator.sess.run(tf.global_variables_initializer())
+
         RemotePPOEvaluator = ray.remote(
             num_cpus=self.config["num_cpus_per_worker"],
             num_gpus=self.config["num_gpus_per_worker"])(PPOEvaluator)
@@ -156,7 +175,7 @@ def standardized(value):
 
         use_gae = self.config["use_gae"]
         dummy = np.zeros_like(samples["advantages"])
-        tuples_per_device = self.local_evaluator.par_opt.load_data(
+        tuples_per_device = self.par_opt.load_data(
             self.local_evaluator.sess,
             [samples["obs"],
              samples["value_targets"] if use_gae else dummy,
@@ -178,7 +197,7 @@ def standardized(value):
                     i == 0 and self.iteration == 0 and
                     batch_index == config["full_trace_nth_sgd_batch"])
                 batch_loss, batch_policy_graph, batch_vf_loss, batch_kl, \
-                    batch_entropy = self.local_evaluator.par_opt.optimize(
+                    batch_entropy = self.par_opt.optimize(
                         self.local_evaluator.sess,
                         permutation[batch_index] * model.per_device_batch_size,
                         extra_ops=list(self.local_evaluator.extra_ops.values()),

python/ray/rllib/ppo/ppo_evaluator.py

@@ -89,32 +89,11 @@ def __init__(self, registry, env_creator, config, logdir, is_remote):
             assert self.batch_size % len(devices) == 0
             self.per_device_batch_size = int(self.batch_size / len(devices))
 
-        def build_loss(obs, vtargets, advs, acts, plog, pvf_preds):
-            return ProximalPolicyGraph(
-                self.env.observation_space, self.env.action_space,
-                obs, vtargets, advs, acts, plog, pvf_preds, self.logit_dim,
-                self.kl_coeff, self.distribution_class, self.config,
-                self.sess, self.registry)
-
-        # TEMPORARY
-        main_thread_scope = tf.get_variable_scope()
-        with tf.variable_scope(main_thread_scope, reuse=tf.AUTO_REUSE):
-            inputs = [
-                self.observations, self.value_targets, self.advantages,
-                self.actions, self.prev_logits, self.prev_vf_preds
-            ]
-            self.common_policy = build_loss(*inputs)
-            self.par_opt = LocalSyncParallelOptimizer(
-                tf.train.AdamOptimizer(self.config["sgd_stepsize"]),
-                self.devices,
-                inputs,
-                self.per_device_batch_size,
-                build_loss,
-                self.logdir)
-
-        # Metric ops
-        self.extra_ops = self.init_extra_ops(
-            self.par_opt.get_device_losses())
+        self.inputs = [
+            self.observations, self.value_targets, self.advantages,
+            self.actions, self.prev_logits, self.prev_vf_preds
+        ]
+        self.common_policy = self.build_loss(*self.inputs)
 
         # References to the model weights
         self.variables = ray.experimental.TensorFlowVariables(
@@ -127,34 +106,39 @@ def build_loss(obs, vtargets, advs, acts, plog, pvf_preds):
         self.sampler = SyncSampler(
             self.env, self.common_policy, self.obs_filter,
             self.config["horizon"], self.config["horizon"])
-        self.sess.run(tf.global_variables_initializer())
 
     def compute_gradients(self, samples):
         raise NotImplementedError
 
     def apply_gradients(self, grads):
         raise NotImplementedError
 
+    def build_loss(self, obs, vtargets, advs, acts, plog, pvf_preds):
+        return ProximalPolicyGraph(
+            self.env.observation_space, self.env.action_space,
+            obs, vtargets, advs, acts, plog, pvf_preds, self.logit_dim,
+            self.kl_coeff, self.distribution_class, self.config,
+            self.sess, self.registry)
+
     def init_extra_ops(self, device_losses):
-        tower_avg_ops = OrderedDict()
+        self.extra_ops = OrderedDict()
         with tf.name_scope("test_outputs"):
             policies = device_losses
-            tower_avg_ops["loss"] = tf.reduce_mean(
+            self.extra_ops["loss"] = tf.reduce_mean(
                 tf.stack(values=[
                     policy.loss for policy in policies]), 0)
-            tower_avg_ops["policy_loss"] = tf.reduce_mean(
+            self.extra_ops["policy_loss"] = tf.reduce_mean(
                 tf.stack(values=[
                     policy.mean_policy_loss for policy in policies]), 0)
-            tower_avg_ops["vf_loss"] = tf.reduce_mean(
+            self.extra_ops["vf_loss"] = tf.reduce_mean(
                 tf.stack(values=[
                     policy.mean_vf_loss for policy in policies]), 0)
-            tower_avg_ops["kl"] = tf.reduce_mean(
+            self.extra_ops["kl"] = tf.reduce_mean(
                 tf.stack(values=[
                     policy.mean_kl for policy in policies]), 0)
-            tower_avg_ops["entropy"] = tf.reduce_mean(
+            self.extra_ops["entropy"] = tf.reduce_mean(
                 tf.stack(values=[
                     policy.mean_entropy for policy in policies]), 0)
-        return tower_avg_ops
 
     def save(self):
         filters = self.get_filters(flush_after=True)