[rllib] Refactor rllib to have a common sample collection pathway

python/ray/rllib/__init__.py

@@ -6,6 +6,11 @@
 # This file is imported from the tune module in order to register RLlib agents.
 from ray.tune.registry import register_trainable
 
+from ray.rllib.utils.policy_graph import PolicyGraph
+from ray.rllib.utils.tf_policy_graph import TFPolicyGraph
+from ray.rllib.utils.common_policy_evaluator import CommonPolicyEvaluator
+from ray.rllib.optimizers.sample_batch import SampleBatch
+
 
 def _register_all():
     for key in ["PPO", "ES", "DQN", "APEX", "A3C", "BC", "PG", "DDPG",
@@ -16,3 +21,7 @@ def _register_all():
 
 
 _register_all()
+
+__all__ = [
+    "PolicyGraph", "TFPolicyGraph", "CommonPolicyEvaluator", "SampleBatch"
+]

python/ray/rllib/a3c/a3c.py

@@ -2,22 +2,21 @@
 from __future__ import division
 from __future__ import print_function
 
-import numpy as np
 import pickle
 import os
 
 import ray
 from ray.rllib.agent import Agent
 from ray.rllib.optimizers import AsyncOptimizer
 from ray.rllib.utils import FilterManager
-from ray.rllib.a3c.a3c_evaluator import A3CEvaluator, RemoteA3CEvaluator, \
-    GPURemoteA3CEvaluator
-from ray.tune.result import TrainingResult
+from ray.rllib.utils.common_policy_evaluator import CommonPolicyEvaluator, \
+    collect_metrics
+from ray.rllib.a3c.common import get_policy_cls
 from ray.tune.trial import Resources
 
 DEFAULT_CONFIG = {
     # Number of workers (excluding master)
-    "num_workers": 4,
+    "num_workers": 2,
     # Size of rollout batch
     "batch_size": 10,
     # Use LSTM model - only applicable for image states
@@ -42,6 +41,8 @@
     "entropy_coeff": -0.01,
     # Whether to place workers on GPUs
     "use_gpu_for_workers": False,
+    # Whether to emit extra summary stats
+    "summarize": False,
     # Model and preprocessor options
     "model": {
         # (Image statespace) - Converts image to Channels = 1
@@ -78,56 +79,48 @@ def default_resource_request(cls, config):
             extra_gpu=cf["use_gpu_for_workers"] and cf["num_workers"] or 0)
 
     def _init(self):
-        self.local_evaluator = A3CEvaluator(
-            self.registry,
-            self.env_creator,
-            self.config,
-            self.logdir,
-            start_sampler=False)
-        if self.config["use_gpu_for_workers"]:
-            remote_cls = GPURemoteA3CEvaluator
+        self.policy_cls = get_policy_cls(self.config)
+
+        if self.config["use_pytorch"]:
+            session_creator = None
         else:
-            remote_cls = RemoteA3CEvaluator
+            import tensorflow as tf
+
+            def session_creator():
+                return tf.Session(
+                    config=tf.ConfigProto(
+                        intra_op_parallelism_threads=1,
+                        inter_op_parallelism_threads=1,
+                        gpu_options=tf.GPUOptions(allow_growth=True)))
+
+        remote_cls = CommonPolicyEvaluator.as_remote(
+            num_gpus=1 if self.config["use_gpu_for_workers"] else 0)
+        self.local_evaluator = CommonPolicyEvaluator(
+            self.env_creator, self.policy_cls,
+            batch_steps=self.config["batch_size"],
+            batch_mode="truncate_episodes",
+            tf_session_creator=session_creator,
+            registry=self.registry, env_config=self.config["env_config"],
+            model_config=self.config["model"], policy_config=self.config)
         self.remote_evaluators = [
-            remote_cls.remote(self.registry, self.env_creator, self.config,
-                              self.logdir)
-            for i in range(self.config["num_workers"])
-        ]
-        self.optimizer = AsyncOptimizer(self.config["optimizer"],
-                                        self.local_evaluator,
-                                        self.remote_evaluators)
+            remote_cls.remote(
+                self.env_creator, self.policy_cls,
+                batch_steps=self.config["batch_size"],
+                batch_mode="truncate_episodes", sample_async=True,
+                tf_session_creator=session_creator,
+                registry=self.registry, env_config=self.config["env_config"],
+                model_config=self.config["model"], policy_config=self.config)
+            for i in range(self.config["num_workers"])]
+
+        self.optimizer = AsyncOptimizer(
+            self.config["optimizer"], self.local_evaluator,
+            self.remote_evaluators)
 
     def _train(self):
         self.optimizer.step()
-        FilterManager.synchronize(self.local_evaluator.filters,
-                                  self.remote_evaluators)
-        res = self._fetch_metrics_from_remote_evaluators()
-        return res
-
-    def _fetch_metrics_from_remote_evaluators(self):
-        episode_rewards = []
-        episode_lengths = []
-        metric_lists = [
-            a.get_completed_rollout_metrics.remote()
-            for a in self.remote_evaluators
-        ]
-        for metrics in metric_lists:
-            for episode in ray.get(metrics):
-                episode_lengths.append(episode.episode_length)
-                episode_rewards.append(episode.episode_reward)
-        avg_reward = (np.mean(episode_rewards)
-                      if episode_rewards else float('nan'))
-        avg_length = (np.mean(episode_lengths)
-                      if episode_lengths else float('nan'))
-        timesteps = np.sum(episode_lengths) if episode_lengths else 0
-
-        result = TrainingResult(
-            episode_reward_mean=avg_reward,
-            episode_len_mean=avg_length,
-            timesteps_this_iter=timesteps,
-            info={})
-
-        return result
+        FilterManager.synchronize(
+            self.local_evaluator.filters, self.remote_evaluators)
+        return collect_metrics(self.local_evaluator, self.remote_evaluators)
 
     def _stop(self):
         # workaround for https://github.com/ray-project/ray/issues/1516
@@ -154,7 +147,10 @@ def _restore(self, checkpoint_path):
         ])
         self.local_evaluator.restore(extra_data["local_state"])
 
-    def compute_action(self, observation):
+    def compute_action(self, observation, state=None):
+        if state is None:
+            state = []
         obs = self.local_evaluator.obs_filter(observation, update=False)
-        action, info = self.local_evaluator.policy.compute(obs)
-        return action
+        return self.local_evaluator.for_policy(
+            lambda p: p.compute_single_action(
+                obs, state, is_training=False)[0])

python/ray/rllib/a3c/a3c_tf_policy.py

@@ -0,0 +1,102 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import tensorflow as tf
+import gym
+from ray.rllib.utils.error import UnsupportedSpaceException
+from ray.rllib.utils.process_rollout import compute_advantages
+from ray.rllib.utils.tf_policy_graph import TFPolicyGraph
+
+
+class A3CTFPolicyGraph(TFPolicyGraph):
+    """The TF policy base class."""
+
+    def __init__(self, ob_space, action_space, registry, config):
+        self.registry = registry
+        self.local_steps = 0
+        self.config = config
+        self.summarize = config.get("summarize")
+
+        self._setup_graph(ob_space, action_space)
+        assert all(hasattr(self, attr)
+                   for attr in ["vf", "logits", "x", "var_list"])
+        print("Setting up loss")
+        self.setup_loss(action_space)
+        self.is_training = tf.placeholder_with_default(True, ())
+        self.sess = tf.get_default_session()
+
+        TFPolicyGraph.__init__(
+            self, self.sess, obs_input=self.x,
+            action_sampler=self.action_dist.sample(), loss=self.loss,
+            loss_inputs=self.loss_in, is_training=self.is_training,
+            state_inputs=self.state_in, state_outputs=self.state_out)
+
+        self.sess.run(tf.global_variables_initializer())
+
+        if self.summarize:
+            bs = tf.to_float(tf.shape(self.x)[0])
+            tf.summary.scalar("model/policy_graph", self.pi_loss / bs)
+            tf.summary.scalar("model/value_loss", self.vf_loss / bs)
+            tf.summary.scalar("model/entropy", self.entropy / bs)
+            tf.summary.scalar("model/grad_gnorm", tf.global_norm(self._grads))
+            tf.summary.scalar("model/var_gnorm", tf.global_norm(self.var_list))
+            self.summary_op = tf.summary.merge_all()
+
+    def _setup_graph(self, ob_space, ac_space):
+        raise NotImplementedError
+
+    def setup_loss(self, action_space):
+        if isinstance(action_space, gym.spaces.Box):
+            ac_size = action_space.shape[0]
+            self.ac = tf.placeholder(tf.float32, [None, ac_size], name="ac")
+        elif isinstance(action_space, gym.spaces.Discrete):
+            self.ac = tf.placeholder(tf.int64, [None], name="ac")
+        else:
+            raise UnsupportedSpaceException(
+                "Action space {} is not supported for A3C.".format(
+                    action_space))
+        self.adv = tf.placeholder(tf.float32, [None], name="adv")
+        self.r = tf.placeholder(tf.float32, [None], name="r")
+
+        log_prob = self.action_dist.logp(self.ac)
+
+        # The "policy gradients" loss: its derivative is precisely the policy
+        # gradient. Notice that self.ac is a placeholder that is provided
+        # externally. adv will contain the advantages, as calculated in
+        # compute_advantages.
+        self.pi_loss = - tf.reduce_sum(log_prob * self.adv)
+
+        delta = self.vf - self.r
+        self.vf_loss = 0.5 * tf.reduce_sum(tf.square(delta))
+        self.entropy = tf.reduce_sum(self.action_dist.entropy())
+        self.loss = (self.pi_loss +
+                     self.vf_loss * self.config["vf_loss_coeff"] +
+                     self.entropy * self.config["entropy_coeff"])
+
+    def optimizer(self):
+        return tf.train.AdamOptimizer(self.config["lr"])
+
+    def gradients(self, optimizer):
+        grads = tf.gradients(self.loss, self.var_list)
+        self.grads, _ = tf.clip_by_global_norm(grads, self.config["grad_clip"])
+        clipped_grads = list(zip(self.grads, self.var_list))
+        return clipped_grads
+
+    def extra_compute_grad_fetches(self):
+        if self.summarize:
+            return {"summary": self.summary_op}
+        else:
+            return {}
+
+    def postprocess_trajectory(self, sample_batch, other_agent_batches=None):
+        completed = sample_batch["dones"][-1]
+        if completed:
+            last_r = 0.0
+        else:
+            next_state = []
+            for i in range(len(self.state_in)):
+                next_state.append([sample_batch["state_out_{}".format(i)][-1]])
+            last_r = self.value(sample_batch["new_obs"][-1], *next_state)
+        return compute_advantages(
+            sample_batch, last_r, self.config["gamma"], self.config["lambda"])