[rllib] Some API cleanups and documentation improvements (#4409)

ericl · web-flow · commit 4b8b70356147 · 2019-03-21T21:34:22.000-07:00
diff --git a/doc/source/rllib-examples.rst b/doc/source/rllib-examples.rst
@@ -26,8 +26,8 @@ Training Workflows
 Custom Envs and Models
 ----------------------
 
-- `Registering a custom env <https://github.com/ray-project/ray/blob/master/python/ray/rllib/examples/custom_env.py>`__:
-   Example of defining and registering a gym env for use with RLlib.
+- `Registering a custom env and model <https://github.com/ray-project/ray/blob/master/python/ray/rllib/examples/custom_env.py>`__:
+   Example of defining and registering a gym env and model for use with RLlib.
 - `Registering a custom model with supervised loss <https://github.com/ray-project/ray/blob/master/python/ray/rllib/examples/custom_loss.py>`__:
    Example of defining and registering a custom model with a supervised loss.
 - `Subprocess environment <https://github.com/ray-project/ray/blob/master/python/ray/rllib/tests/test_env_with_subprocess.py>`__:
diff --git a/doc/source/rllib-offline.rst b/doc/source/rllib-offline.rst
@@ -157,8 +157,8 @@ You can configure experience input for an agent using the following options:
 
 .. literalinclude:: ../../python/ray/rllib/agents/agent.py
    :language: python
-   :start-after: __sphinx_doc_input_begin__
-   :end-before: __sphinx_doc_input_end__
+   :start-after: === Offline Datasets ===
+   :end-before: Specify where experiences should be saved
 
 The interface for a custom input reader is as follows:
 
@@ -172,8 +172,8 @@ You can configure experience output for an agent using the following options:
 
 .. literalinclude:: ../../python/ray/rllib/agents/agent.py
    :language: python
-   :start-after: __sphinx_doc_output_begin__
-   :end-before: __sphinx_doc_output_end__
+   :start-after: shuffle_buffer_size
+   :end-before: === Multiagent ===
 
 The interface for a custom output writer is as follows:
 
diff --git a/doc/source/rllib-training.rst b/doc/source/rllib-training.rst
@@ -37,7 +37,7 @@ The ``rllib train`` command (same as the ``train.py`` script in the repo) has a
 The most important options are for choosing the environment
 with ``--env`` (any OpenAI gym environment including ones registered by the user
 can be used) and for choosing the algorithm with ``--run``
-(available options are ``PPO``, ``PG``, ``A2C``, ``A3C``, ``IMPALA``, ``ES``, ``DDPG``, ``DQN``, ``APEX``, and ``APEX_DDPG``).
+(available options are ``PPO``, ``PG``, ``A2C``, ``A3C``, ``IMPALA``, ``ES``, ``DDPG``, ``DQN``, ``MARWIL``, ``APEX``, and ``APEX_DDPG``).
 
 Evaluating Trained Agents
 ~~~~~~~~~~~~~~~~~~~~~~~~~
diff --git a/doc/source/rllib.rst b/doc/source/rllib.rst
@@ -1,7 +1,7 @@
 RLlib: Scalable Reinforcement Learning
 ======================================
 
-RLlib is an open-source library for reinforcement learning that offers both a collection of reference algorithms and scalable primitives for composing new ones.
+RLlib is an open-source library for reinforcement learning that offers both a unified API for a variety of applications and high scalability via distributed eager execution.
 
 .. image:: rllib-stack.svg
 
diff --git a/python/ray/rllib/README.md b/python/ray/rllib/README.md
@@ -1,7 +1,7 @@
 RLlib: Scalable Reinforcement Learning
 ======================================
 
-RLlib is an open-source library for reinforcement learning that offers both a collection of reference algorithms and scalable primitives for composing new ones.
+RLlib is an open-source library for reinforcement learning that offers both a unified API for a variety of applications and high scalability via distributed eager execution.
 
 For an overview of RLlib, see the [documentation](http://ray.readthedocs.io/en/latest/rllib.html).
 
diff --git a/python/ray/rllib/agents/a3c/a2c.py b/python/ray/rllib/agents/a3c/a2c.py
@@ -25,6 +25,6 @@ class A2CAgent(A3CAgent):
 
     @override(A3CAgent)
     def _make_optimizer(self):
-        return SyncSamplesOptimizer(self.local_evaluator,
-                                    self.remote_evaluators,
-                                    self.config["optimizer"])
+        return SyncSamplesOptimizer(
+            self.local_evaluator, self.remote_evaluators,
+            {"train_batch_size": self.config["train_batch_size"]})
diff --git a/python/ray/rllib/agents/a3c/a3c.py b/python/ray/rllib/agents/a3c/a3c.py
@@ -69,8 +69,7 @@ def _train(self):
         start = time.time()
         while time.time() - start < self.config["min_iter_time_s"]:
             self.optimizer.step()
-        result = self.optimizer.collect_metrics(
-            self.config["collect_metrics_timeout"])
+        result = self.collect_metrics()
         result.update(timesteps_this_iter=self.optimizer.num_steps_sampled -
                       prev_steps)
         return result
diff --git a/python/ray/rllib/agents/a3c/a3c_tf_policy_graph.py b/python/ray/rllib/agents/a3c/a3c_tf_policy_graph.py
@@ -146,8 +146,8 @@ def postprocess_trajectory(self,
                                   self.config["lambda"])
 
     @override(TFPolicyGraph)
-    def gradients(self, optimizer):
-        grads = tf.gradients(self._loss, self.var_list)
+    def gradients(self, optimizer, loss):
+        grads = tf.gradients(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
diff --git a/python/ray/rllib/agents/agent.py b/python/ray/rllib/agents/agent.py
@@ -99,7 +99,9 @@
     # === Execution ===
     # Number of environments to evaluate vectorwise per worker.
     "num_envs_per_worker": 1,
-    # Default sample batch size
+    # Default sample batch size (unroll length). Batches of this size are
+    # collected from workers until train_batch_size is met. When using
+    # multiple envs per worker, this is multiplied by num_envs_per_worker.
     "sample_batch_size": 200,
     # Training batch size, if applicable. Should be >= sample_batch_size.
     # Samples batches will be concatenated together to this size for training.
@@ -137,6 +139,8 @@
     "compress_observations": False,
     # Drop metric batches from unresponsive workers after this many seconds
     "collect_metrics_timeout": 180,
+    # Smooth metrics over this many episodes.
+    "metrics_smoothing_episodes": 100,
     # If using num_envs_per_worker > 1, whether to create those new envs in
     # remote processes instead of in the same worker. This adds overheads, but
     # can make sense if your envs are very CPU intensive (e.g., for StarCraft).
@@ -146,7 +150,6 @@
     "async_remote_worker_envs": False,
 
     # === Offline Datasets ===
-    # __sphinx_doc_input_begin__
     # Specify how to generate experiences:
     #  - "sampler": generate experiences via online simulation (default)
     #  - a local directory or file glob expression (e.g., "/tmp/*.json")
@@ -172,8 +175,6 @@
     # of this number of batches. Use this if the input data is not in random
     # enough order. Input is delayed until the shuffle buffer is filled.
     "shuffle_buffer_size": 0,
-    # __sphinx_doc_input_end__
-    # __sphinx_doc_output_begin__
     # Specify where experiences should be saved:
     #  - None: don't save any experiences
     #  - "logdir" to save to the agent log dir
@@ -184,7 +185,6 @@
     "output_compress_columns": ["obs", "new_obs"],
     # Max output file size before rolling over to a new file.
     "output_max_file_size": 64 * 1024 * 1024,
-    # __sphinx_doc_output_end__
 
     # === Multiagent ===
     "multiagent": {
@@ -559,6 +559,17 @@ def export_policy_checkpoint(self,
         self.local_evaluator.export_policy_checkpoint(
             export_dir, filename_prefix, policy_id)
 
+    @DeveloperAPI
+    def collect_metrics(self, selected_evaluators=None):
+        """Collects metrics from the remote evaluators of this agent.
+
+        This is the same data as returned by a call to train().
+        """
+        return self.optimizer.collect_metrics(
+            self.config["collect_metrics_timeout"],
+            min_history=self.config["metrics_smoothing_episodes"],
+            selected_evaluators=selected_evaluators)
+
     @classmethod
     def resource_help(cls, config):
         return ("\n\nYou can adjust the resource requests of RLlib agents by "
diff --git a/python/ray/rllib/agents/ddpg/apex.py b/python/ray/rllib/agents/ddpg/apex.py
@@ -23,7 +23,6 @@
         "learning_starts": 50000,
         "train_batch_size": 512,
         "sample_batch_size": 50,
-        "max_weight_sync_delay": 400,
         "target_network_update_freq": 500000,
         "timesteps_per_iteration": 25000,
         "per_worker_exploration": True,
diff --git a/python/ray/rllib/agents/ddpg/ddpg.py b/python/ray/rllib/agents/ddpg/ddpg.py
@@ -37,11 +37,11 @@
     "evaluation_num_episodes": 10,
 
     # === Model ===
-    # Hidden layer sizes of the policy network
+    # Postprocess the policy network model output with these hidden layers
     "actor_hiddens": [64, 64],
     # Hidden layers activation of the policy network
     "actor_hidden_activation": "relu",
-    # Hidden layer sizes of the critic network
+    # Postprocess the critic network model output with these hidden layers
     "critic_hiddens": [64, 64],
     # Hidden layers activation of the critic network
     "critic_hidden_activation": "relu",
diff --git a/python/ray/rllib/agents/ddpg/ddpg_policy_graph.py b/python/ray/rllib/agents/ddpg/ddpg_policy_graph.py
@@ -408,7 +408,7 @@ def optimizer(self):
         return tf.train.AdamOptimizer(learning_rate=self.config["lr"])
 
     @override(TFPolicyGraph)
-    def gradients(self, optimizer):
+    def gradients(self, optimizer, loss):
         if self.config["grad_norm_clipping"] is not None:
             actor_grads_and_vars = _minimize_and_clip(
                 optimizer,
diff --git a/python/ray/rllib/agents/dqn/dqn.py b/python/ray/rllib/agents/dqn/dqn.py
@@ -41,7 +41,8 @@
     "dueling": True,
     # Whether to use double dqn
     "double_q": True,
-    # Hidden layer sizes of the state and action value networks
+    # Postprocess model outputs with these hidden layers to compute the
+    # state and action values. See also the model config in catalog.py.
     "hiddens": [256],
     # N-step Q learning
     "n_step": 1,
@@ -69,7 +70,7 @@
     "exploration_final_eps": 0.02,
     # Update the target network every `target_network_update_freq` steps.
     "target_network_update_freq": 500,
-    # Use softmax for sampling actions.
+    # Use softmax for sampling actions. Required for off policy estimation.
     "soft_q": False,
     # Softmax temperature. Q values are divided by this value prior to softmax.
     # Softmax approaches argmax as the temperature drops to zero.
@@ -260,13 +261,11 @@ def _train(self):
 
         if self.config["per_worker_exploration"]:
             # Only collect metrics from the third of workers with lowest eps
-            result = self.optimizer.collect_metrics(
-                timeout_seconds=self.config["collect_metrics_timeout"],
+            result = self.collect_metrics(
                 selected_evaluators=self.remote_evaluators[
                     -len(self.remote_evaluators) // 3:])
         else:
-            result = self.optimizer.collect_metrics(
-                timeout_seconds=self.config["collect_metrics_timeout"])
+            result = self.collect_metrics()
 
         result.update(
             timesteps_this_iter=self.global_timestep - start_timestep,
diff --git a/python/ray/rllib/agents/dqn/dqn_policy_graph.py b/python/ray/rllib/agents/dqn/dqn_policy_graph.py
@@ -423,16 +423,16 @@ def optimizer(self):
             epsilon=self.config["adam_epsilon"])
 
     @override(TFPolicyGraph)
-    def gradients(self, optimizer):
+    def gradients(self, optimizer, loss):
         if self.config["grad_norm_clipping"] is not None:
             grads_and_vars = _minimize_and_clip(
                 optimizer,
-                self._loss,
+                loss,
                 var_list=self.q_func_vars,
                 clip_val=self.config["grad_norm_clipping"])
         else:
             grads_and_vars = optimizer.compute_gradients(
-                self.loss.loss, var_list=self.q_func_vars)
+                loss, var_list=self.q_func_vars)
         grads_and_vars = [(g, v) for (g, v) in grads_and_vars if g is not None]
         return grads_and_vars
 
diff --git a/python/ray/rllib/agents/impala/impala.py b/python/ray/rllib/agents/impala/impala.py
@@ -121,8 +121,7 @@ def _train(self):
         while (time.time() - start < self.config["min_iter_time_s"]
                or self.optimizer.num_steps_sampled == prev_steps):
             self.optimizer.step()
-        result = self.optimizer.collect_metrics(
-            self.config["collect_metrics_timeout"])
+        result = self.collect_metrics()
         result.update(timesteps_this_iter=self.optimizer.num_steps_sampled -
                       prev_steps)
         return result
diff --git a/python/ray/rllib/agents/impala/vtrace_policy_graph.py b/python/ray/rllib/agents/impala/vtrace_policy_graph.py
@@ -327,8 +327,8 @@ def optimizer(self):
                                              self.config["epsilon"])
 
     @override(TFPolicyGraph)
-    def gradients(self, optimizer):
-        grads = tf.gradients(self._loss, self.var_list)
+    def gradients(self, optimizer, loss):
+        grads = tf.gradients(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
diff --git a/python/ray/rllib/agents/marwil/marwil.py b/python/ray/rllib/agents/marwil/marwil.py
@@ -63,8 +63,7 @@ def _init(self):
     def _train(self):
         prev_steps = self.optimizer.num_steps_sampled
         fetches = self.optimizer.step()
-        res = self.optimizer.collect_metrics(
-            self.config["collect_metrics_timeout"])
+        res = self.collect_metrics()
         res.update(
             timesteps_this_iter=self.optimizer.num_steps_sampled - prev_steps,
             info=dict(fetches, **res.get("info", {})))
diff --git a/python/ray/rllib/agents/pg/pg.py b/python/ray/rllib/agents/pg/pg.py
@@ -55,8 +55,7 @@ def _init(self):
     def _train(self):
         prev_steps = self.optimizer.num_steps_sampled
         self.optimizer.step()
-        result = self.optimizer.collect_metrics(
-            self.config["collect_metrics_timeout"])
+        result = self.collect_metrics()
         result.update(timesteps_this_iter=self.optimizer.num_steps_sampled -
                       prev_steps)
         return result
diff --git a/python/ray/rllib/agents/ppo/appo_policy_graph.py b/python/ray/rllib/agents/ppo/appo_policy_graph.py
@@ -427,8 +427,8 @@ def optimizer(self):
                                              self.config["momentum"],
                                              self.config["epsilon"])
 
-    def gradients(self, optimizer):
-        grads = tf.gradients(self.loss.total_loss, self.var_list)
+    def gradients(self, optimizer, loss):
+        grads = tf.gradients(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
diff --git a/python/ray/rllib/agents/ppo/ppo.py b/python/ray/rllib/agents/ppo/ppo.py
@@ -124,8 +124,7 @@ def update(pi, pi_id):
 
             # multi-agent
             self.local_evaluator.foreach_trainable_policy(update)
-        res = self.optimizer.collect_metrics(
-            self.config["collect_metrics_timeout"])
+        res = self.collect_metrics()
         res.update(
             timesteps_this_iter=self.optimizer.num_steps_sampled - prev_steps,
             info=dict(fetches, **res.get("info", {})))
diff --git a/python/ray/rllib/agents/ppo/ppo_policy_graph.py b/python/ray/rllib/agents/ppo/ppo_policy_graph.py
@@ -305,18 +305,18 @@ def postprocess_trajectory(self,
         return batch
 
     @override(TFPolicyGraph)
-    def gradients(self, optimizer):
+    def gradients(self, optimizer, loss):
         if self.config["grad_clip"] is not None:
             self.var_list = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES,
                                               tf.get_variable_scope().name)
-            grads = tf.gradients(self._loss, self.var_list)
+            grads = tf.gradients(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
         else:
             return optimizer.compute_gradients(
-                self._loss, colocate_gradients_with_ops=True)
+                loss, colocate_gradients_with_ops=True)
 
     @override(PolicyGraph)
     def get_initial_state(self):
diff --git a/python/ray/rllib/agents/qmix/apex.py b/python/ray/rllib/agents/qmix/apex.py
@@ -26,7 +26,6 @@
         "learning_starts": 50000,
         "train_batch_size": 512,
         "sample_batch_size": 50,
-        "max_weight_sync_delay": 400,
         "target_network_update_freq": 500000,
         "timesteps_per_iteration": 25000,
         "per_worker_exploration": True,
diff --git a/python/ray/rllib/evaluation/tf_policy_graph.py b/python/ray/rllib/evaluation/tf_policy_graph.py
@@ -128,9 +128,10 @@ def __init__(self,
             self._stats_fetches = {}
 
         self._optimizer = self.optimizer()
-        self._grads_and_vars = [(g, v)
-                                for (g, v) in self.gradients(self._optimizer)
-                                if g is not None]
+        self._grads_and_vars = [
+            (g, v) for (g, v) in self.gradients(self._optimizer, self._loss)
+            if g is not None
+        ]
         self._grads = [g for (g, v) in self._grads_and_vars]
         self._variables = ray.experimental.tf_utils.TensorFlowVariables(
             self._loss, self._sess)
@@ -145,10 +146,8 @@ def __init__(self,
             logger.debug("Update ops to run on apply gradient: {}".format(
                 self._update_ops))
         with tf.control_dependencies(self._update_ops):
-            # specify global_step for TD3 which needs to count the num updates
-            self._apply_op = self._optimizer.apply_gradients(
-                self._grads_and_vars,
-                global_step=tf.train.get_or_create_global_step())
+            self._apply_op = self.build_apply_op(self._optimizer,
+                                                 self._grads_and_vars)
 
         if len(self._state_inputs) != len(self._state_outputs):
             raise ValueError(
@@ -281,9 +280,18 @@ def optimizer(self):
         return tf.train.AdamOptimizer()
 
     @DeveloperAPI
-    def gradients(self, optimizer):
+    def gradients(self, optimizer, loss):
         """Override for custom gradient computation."""
-        return optimizer.compute_gradients(self._loss)
+        return optimizer.compute_gradients(loss)
+
+    @DeveloperAPI
+    def build_apply_op(self, optimizer, grads_and_vars):
+        """Override for custom gradient apply computation."""
+
+        # specify global_step for TD3 which needs to count the num updates
+        return optimizer.apply_gradients(
+            self._grads_and_vars,
+            global_step=tf.train.get_or_create_global_step())
 
     @DeveloperAPI
     def _get_is_training_placeholder(self):
diff --git a/python/ray/rllib/examples/custom_env.py b/python/ray/rllib/examples/custom_env.py
diff --git a/python/ray/rllib/optimizers/multi_gpu_impl.py b/python/ray/rllib/optimizers/multi_gpu_impl.py
diff --git a/python/ray/rllib/tuned_examples/pendulum-ppo.yaml b/python/ray/rllib/tuned_examples/pendulum-ppo.yaml
