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

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).
 

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"]})

python/ray/rllib/agents/a3c/a3c.py

@@ -69,8 +69,7 @@ class A3CAgent(Agent):
         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

python/ray/rllib/agents/a3c/a3c_tf_policy_graph.py

@@ -146,8 +146,8 @@ class A3CPolicyGraph(LearningRateSchedule, TFPolicyGraph):
                                   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

python/ray/rllib/agents/agent.py

@@ -99,7 +99,9 @@ COMMON_CONFIG = {
     # === 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 @@ COMMON_CONFIG = {
     "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 @@ COMMON_CONFIG = {
     "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 @@ COMMON_CONFIG = {
     # 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 @@ COMMON_CONFIG = {
     "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 @@ class Agent(Trainable):
         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 "

python/ray/rllib/agents/ddpg/apex.py

@@ -23,7 +23,6 @@ APEX_DDPG_DEFAULT_CONFIG = merge_dicts(
         "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,

python/ray/rllib/agents/ddpg/ddpg.py

@@ -37,11 +37,11 @@ DEFAULT_CONFIG = with_common_config({
     "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",

python/ray/rllib/agents/ddpg/ddpg_policy_graph.py

@@ -408,7 +408,7 @@ class DDPGPolicyGraph(TFPolicyGraph):
         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,

python/ray/rllib/agents/dqn/dqn.py

@@ -41,7 +41,8 @@ DEFAULT_CONFIG = with_common_config({
     "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 @@ DEFAULT_CONFIG = with_common_config({
     "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 @@ class DQNAgent(Agent):
 
         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,

python/ray/rllib/agents/dqn/dqn_policy_graph.py

@@ -423,16 +423,16 @@ class DQNPolicyGraph(TFPolicyGraph):
             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
 

python/ray/rllib/agents/impala/impala.py

@@ -121,8 +121,7 @@ class ImpalaAgent(Agent):
         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

python/ray/rllib/agents/impala/vtrace_policy_graph.py

@@ -327,8 +327,8 @@ class VTracePolicyGraph(LearningRateSchedule, TFPolicyGraph):
                                              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

python/ray/rllib/agents/marwil/marwil.py

@@ -63,8 +63,7 @@ class MARWILAgent(Agent):
     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", {})))

python/ray/rllib/agents/pg/pg.py

@@ -55,8 +55,7 @@ class PGAgent(Agent):
     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

python/ray/rllib/agents/ppo/appo_policy_graph.py

@@ -427,8 +427,8 @@ class AsyncPPOPolicyGraph(LearningRateSchedule, TFPolicyGraph):
                                              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

python/ray/rllib/agents/ppo/ppo.py

@@ -124,8 +124,7 @@ class PPOAgent(Agent):
 
             # 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", {})))

python/ray/rllib/agents/ppo/ppo_policy_graph.py

@@ -305,18 +305,18 @@ class PPOPolicyGraph(LearningRateSchedule, TFPolicyGraph):
         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):

python/ray/rllib/agents/qmix/apex.py

@@ -26,7 +26,6 @@ APEX_QMIX_DEFAULT_CONFIG = merge_dicts(
         "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,

python/ray/rllib/evaluation/tf_policy_graph.py

@@ -128,9 +128,10 @@ class TFPolicyGraph(PolicyGraph):
             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 @@ class TFPolicyGraph(PolicyGraph):
             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 @@ class TFPolicyGraph(PolicyGraph):
         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):

python/ray/rllib/examples/custom_env.py

@@ -1,7 +1,8 @@
-"""Example of a custom gym environment. Run this for a demo.
+"""Example of a custom gym environment and model. Run this for a demo.
 
 This example shows:
   - using a custom environment
+  - using a custom model
   - using Tune for grid search
 
 You can visualize experiment results in ~/ray_results using TensorBoard.
@@ -13,6 +14,7 @@ from __future__ import print_function
 
 import numpy as np
 import gym
+from ray.rllib.models import FullyConnectedNetwork, Model, ModelCatalog
 from gym.spaces import Discrete, Box
 
 import ray
@@ -45,10 +47,25 @@ class SimpleCorridor(gym.Env):
         return [self.cur_pos], 1 if done else 0, done, {}
 
 
+class CustomModel(Model):
+    """Example of a custom model.
+
+    This model just delegates to the built-in fcnet.
+    """
+
+    def _build_layers_v2(self, input_dict, num_outputs, options):
+        self.obs_in = input_dict["obs"]
+        self.fcnet = FullyConnectedNetwork(input_dict, self.obs_space,
+                                           self.action_space, num_outputs,
+                                           options)
+        return self.fcnet.outputs, self.fcnet.last_layer
+
+
 if __name__ == "__main__":
     # Can also register the env creator function explicitly with:
     # register_env("corridor", lambda config: SimpleCorridor(config))
     ray.init()
+    ModelCatalog.register_custom_model("my_model", CustomModel)
     run_experiments({
         "demo": {
             "run": "PPO",
@@ -57,6 +74,9 @@ if __name__ == "__main__":
                 "timesteps_total": 10000,
             },
             "config": {
+                "model": {
+                    "custom_model": "my_model",
+                },
                 "lr": grid_search([1e-2, 1e-4, 1e-6]),  # try different lrs
                 "num_workers": 1,  # parallelism
                 "env_config": {

python/ray/rllib/optimizers/multi_gpu_impl.py

@@ -262,7 +262,8 @@ class LocalSyncParallelOptimizer(object):
                     current_slice.set_shape(ph.shape)
                     device_input_slices.append(current_slice)
                 graph_obj = self.build_graph(device_input_slices)
-                device_grads = graph_obj.gradients(self.optimizer)
+                device_grads = graph_obj.gradients(self.optimizer,
+                                                   graph_obj._loss)
             return Tower(
                 tf.group(
                     *[batch.initializer for batch in device_input_batches]),

python/ray/rllib/tuned_examples/pendulum-ppo.yaml

@@ -15,4 +15,4 @@ pendulum-ppo:
         model:
             fcnet_hiddens: [64, 64]
         batch_mode: complete_episodes
-        observation_fliter: MeanStdFilter
+        observation_filter: MeanStdFilter