[rllib] Initial work on integrating hyperparameter search tool (#1107)

* clean up train

* update

* update train script

* add tuned examples

* add agent catalog

* add tune lib

* update

* fix

* testS

* remove

* train docs

* comments

* todo

* fix resource parsing

* fix cr test

* add test

* try to fix travis test

python/ray/tune/README.rst

@@ -0,0 +1,54 @@
+Ray.tune: Fast hyperparameter search
+====================================
+
+Using ray.tune with RLlib
+-------------------------
+
+One way to use ray.tune is through RLlib's train.py script. The train.py script
+supports two modes. For example, to run multiple concurrent trials of Pong:
+
+- Inline args: ``./train.py --env=Pong-v0 --alg=PPO --num_trials=8 --stop '{"time_total_s": 3200}' --resources '{"cpu": 8, "gpu": 2}' --config '{"num_workers": 8, "sgd_num_iter": 10}'``
+
+- File-based: ``./train.py -f tune-pong.yaml``
+
+Both delegate scheduling of trials to the ray.tune TrialRunner class.
+Additionally, the file-based mode supports hyper-parameter tuning
+(currently just grid and random search).
+
+To specify search parameters, variables in the `config` section may be set to
+different values for each trial. You can either specify `grid_search: <list>`
+in place of a concrete value to specify a grid search across the list of
+values, or `eval: <str>` for values to be sampled from the given Python
+expression.
+
+.. code:: yaml
+
+    cartpole-ppo:
+        env: CartPole-v0
+        alg: PPO
+        num_trials: 6
+        stop:
+            episode_reward_mean: 200
+            time_total_s: 180
+        resources:
+            cpu: 4
+        config:
+            num_workers: 4
+            num_sgd_iter:
+                grid_search: [1, 4]
+            sgd_batchsize:
+                grid_search: [128, 256, 512]
+            lr:
+                eval: random.uniform(1e-4, 1e-3)
+
+See ray/rllib/tuned_examples for more examples of configs in YAML form.
+
+Using ray.tune to run custom scripts
+------------------------------------
+
+TODO
+
+Using ray.tune as a library
+---------------------------
+
+TODO

python/ray/tune/config_parser.py

@@ -0,0 +1,138 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+
+import argparse
+import json
+import numpy as np
+import os
+import random
+import sys
+
+from ray.tune.trial import Trial, Resources
+
+
+def _resource_json(data):
+    values = json.loads(data)
+    return Resources(values.get('cpu', 0), values.get('gpu', 0))
+
+
+def make_parser(description):
+    """Returns a base argument parser for the ray.tune tool."""
+
+    parser = argparse.ArgumentParser(description=(description))
+
+    parser.add_argument("--alg", default="PPO", type=str,
+                        help="The learning algorithm to train.")
+    parser.add_argument("--stop", default="{}", type=json.loads,
+                        help="The stopping criteria, specified in JSON.")
+    parser.add_argument("--config", default="{}", type=json.loads,
+                        help="The config of the algorithm, specified in JSON.")
+    parser.add_argument("--resources", default='{"cpu": 1}',
+                        type=_resource_json,
+                        help="Amount of resources to allocate per trial.")
+    parser.add_argument("--num_trials", default=1, type=int,
+                        help="Number of trials to evaluate.")
+    parser.add_argument("--local_dir", default="/tmp/ray", type=str,
+                        help="Local dir to save training results to.")
+    parser.add_argument("--upload_dir", default=None, type=str,
+                        help="URI to upload training results to.")
+    parser.add_argument("--checkpoint_freq", default=sys.maxsize, type=int,
+                        help="How many iterations between checkpoints.")
+
+    # TODO(ekl) environments are RL specific
+    parser.add_argument("--env", default=None, type=str,
+                        help="The gym environment to use.")
+
+    return parser
+
+
+def parse_to_trials(config):
+    """Parses a json config to the number of trials specified by the config.
+
+    The input config is a mapping from experiment names to an argument
+    dictionary describing a set of trials. These args include the parser args
+    documented in make_parser().
+    """
+
+    def resolve(agent_cfg, resolved_vars, i):
+        assert type(agent_cfg) == dict
+        cfg = agent_cfg.copy()
+        for p, val in cfg.items():
+            if type(val) == dict and "eval" in val:
+                cfg[p] = eval(val["eval"], {
+                    "random": random,
+                    "np": np,
+                }, {
+                    "_i": i,
+                })
+                resolved_vars[p] = True
+        return cfg, resolved_vars
+
+    def to_argv(config):
+        argv = []
+        for k, v in config.items():
+            argv.append("--{}".format(k))
+            if type(v) is str:
+                argv.append(v)
+            else:
+                argv.append(json.dumps(v))
+        return argv
+
+    def param_str(config, resolved_vars):
+        return "_".join(
+            [k + "=" + str(v) for k, v in sorted(config.items())
+                if resolved_vars.get(k)])
+
+    parser = make_parser("Ray hyperparameter tuning tool")
+    trials = []
+    for experiment_name, exp_cfg in config.items():
+        args = parser.parse_args(to_argv(exp_cfg))
+        grid_search = _GridSearchGenerator(args.config)
+        for i in range(args.num_trials):
+            next_cfg, resolved_vars = grid_search.next()
+            resolved, resolved_vars = resolve(next_cfg, resolved_vars, i)
+            if resolved_vars:
+                agent_id = "{}_{}".format(
+                    i, param_str(resolved, resolved_vars))
+            else:
+                agent_id = str(i)
+            trials.append(Trial(
+                args.env, args.alg, resolved,
+                os.path.join(args.local_dir, experiment_name), agent_id,
+                args.resources, args.stop, args.checkpoint_freq, None,
+                args.upload_dir))
+
+    return trials
+
+
+class _GridSearchGenerator(object):
+    """Generator that implements grid search over a set of value lists."""
+
+    def __init__(self, agent_cfg):
+        self.cfg = agent_cfg
+        self.grid_values = []
+        for p, val in sorted(agent_cfg.items()):
+            if type(val) == dict and "grid_search" in val:
+                assert type(val["grid_search"] == list)
+                self.grid_values.append((p, val["grid_search"]))
+        self.value_indices = [0] * len(self.grid_values)
+
+    def next(self):
+        cfg = self.cfg.copy()
+        resolved_vars = {}
+        for i, (k, values) in enumerate(self.grid_values):
+            idx = self.value_indices[i]
+            cfg[k] = values[idx]
+            resolved_vars[k] = True
+        if self.grid_values:
+            self._increment(0)
+        return cfg, resolved_vars
+
+    def _increment(self, i):
+        self.value_indices[i] += 1
+        if self.value_indices[i] >= len(self.grid_values[i][1]):
+            self.value_indices[i] = 0
+            if i + 1 < len(self.value_indices):
+                self._increment(i + 1)

python/ray/tune/trial.py

@@ -0,0 +1,167 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import sys
+import traceback
+import ray
+
+from collections import namedtuple
+from ray.rllib.agents import get_agent_class
+
+
+# Ray resources required to schedule a Trial
+Resources = namedtuple("Resources", ["cpu", "gpu"])
+
+
+class Trial(object):
+    """A trial object holds the state for one model training run.
+
+    Trials are themselves managed by the TrialRunner class, which implements
+    the event loop for submitting trial runs to a Ray cluster.
+
+    Trials start in the PENDING state, and transition to RUNNING once started.
+    On error it transitions to ERROR, otherwise TERMINATED on success.
+    """
+
+    PENDING = 'PENDING'
+    RUNNING = 'RUNNING'
+    TERMINATED = 'TERMINATED'
+    ERROR = 'ERROR'
+
+    def __init__(
+            self, env_creator, alg, config={}, local_dir='/tmp/ray',
+            agent_id=None, resources=Resources(cpu=1, gpu=0),
+            stopping_criterion={}, checkpoint_freq=sys.maxsize,
+            restore_path=None, upload_dir=None):
+        """Initialize a new trial.
+
+        The args here take the same meaning as the command line flags defined
+        in ray.tune.config_parser.
+        """
+
+        # Immutable config
+        self.env_creator = env_creator
+        if type(env_creator) is str:
+            self.env_name = env_creator
+        else:
+            self.env_name = "custom"
+        self.alg = alg
+        self.config = config
+        self.local_dir = local_dir
+        self.agent_id = agent_id
+        self.resources = resources
+        self.stopping_criterion = stopping_criterion
+        self.checkpoint_freq = checkpoint_freq
+        self.restore_path = restore_path
+        self.upload_dir = upload_dir
+
+        # Local trial state that is updated during the run
+        self.last_result = None
+        self.checkpoint_path = None
+        self.agent = None
+        self.status = Trial.PENDING
+
+    def start(self):
+        """Starts this trial.
+
+        If an error is encountered when starting the trial, an exception will
+        be thrown.
+        """
+
+        self.status = Trial.RUNNING
+        agent_cls = get_agent_class(self.alg)
+        cls = ray.remote(
+            num_cpus=self.resources.cpu, num_gpus=self.resources.gpu)(
+                agent_cls)
+        self.agent = cls.remote(
+            self.env_creator, self.config, self.local_dir, self.upload_dir,
+            agent_id=self.agent_id)
+        if self.restore_path:
+            ray.get(self.agent.restore.remote(self.restore_path))
+
+    def stop(self, error=False):
+        """Stops this trial.
+
+        Stops this trial, releasing all allocating resources. If stopping the
+        trial fails, the run will be marked as terminated in error, but no
+        exception will be thrown.
+
+        Args:
+            error (bool): Whether to mark this trial as terminated in error.
+        """
+
+        if error:
+            self.status = Trial.ERROR
+        else:
+            self.status = Trial.TERMINATED
+
+        try:
+            if self.agent:
+                self.agent.stop.remote()
+                self.agent.__ray_terminate__.remote(
+                    self.agent._ray_actor_id.id())
+        except:
+            print("Error stopping agent:", traceback.format_exc())
+            self.status = Trial.ERROR
+        finally:
+            self.agent = None
+
+    def train_remote(self):
+        """Returns Ray future for one iteration of training."""
+
+        assert self.status == Trial.RUNNING, self.status
+        return self.agent.train.remote()
+
+    def should_stop(self, result):
+        """Whether the given result meets this trial's stopping criteria."""
+
+        for criteria, stop_value in self.stopping_criterion.items():
+            if getattr(result, criteria) >= stop_value:
+                return True
+
+        return False
+
+    def should_checkpoint(self):
+        """Whether this trial is due for checkpointing."""
+
+        if self.checkpoint_freq is None:
+            return False
+
+        return self.last_result.training_iteration % self.checkpoint_freq == 0
+
+    def progress_string(self):
+        """Returns a progress message for printing out to the console."""
+
+        if self.last_result is None:
+            return self.status
+        return '{}, {} s, {} ts, {} itrs, {} rew'.format(
+            self.status,
+            int(self.last_result.time_total_s),
+            int(self.last_result.timesteps_total),
+            self.last_result.training_iteration,
+            round(self.last_result.episode_reward_mean, 1))
+
+    def checkpoint(self):
+        """Synchronously checkpoints the state of this trial.
+
+        TODO(ekl): we should support a PAUSED state based on checkpointing.
+        """
+
+        path = ray.get(self.agent.save.remote())
+        self.checkpoint_path = path
+        print("Saved checkpoint to:", path)
+
+        return path
+
+    def __str__(self):
+        identifier = '{}_{}'.format(self.alg, self.env_name)
+        if self.agent_id:
+            identifier += '_' + self.agent_id
+        return identifier
+
+    def __eq__(self, other):
+        return str(self) == str(other)
+
+    def __hash__(self):
+        return hash(str(self))

python/ray/tune/trial_runner.py

@@ -0,0 +1,182 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import ray
+import time
+import traceback
+
+from ray.tune.trial import Trial, Resources
+
+
+class TrialRunner(object):
+    """A TrialRunner implements the event loop for scheduling trials on Ray.
+
+    Example:
+        runner = TrialRunner()
+        runner.add_trial(Trial(...))
+        runner.add_trial(Trial(...))
+        while not runner.is_finished():
+            runner.step()
+            print(runner.debug_string())
+
+    The main job of TrialRunner is scheduling trials to efficiently use cluster
+    resources, without overloading the cluster.
+
+    While Ray itself provides resource management for tasks and actors, this is
+    not sufficient when scheduling trials that may instantiate multiple actors.
+    This is because if insufficient resources are available, concurrent agents
+    could deadlock waiting for new resources to become available. Furthermore,
+    oversubscribing the cluster could degrade training performance, leading to
+    misleading benchmark results.
+    """
+
+    def __init__(self):
+        """Initializes a new TrialRunner."""
+
+        self._trials = []
+        self._pending = {}
+        self._avail_resources = Resources(cpu=0, gpu=0)
+        self._committed_resources = Resources(cpu=0, gpu=0)
+
+    def is_finished(self):
+        """Returns whether all trials have finished running."""
+
+        for t in self._trials:
+            if t.status in [Trial.PENDING, Trial.RUNNING]:
+                return False
+        return True
+
+    def step(self):
+        """Runs one step of the trial event loop.
+
+        Callers should typically run this method repeatedly in a loop. They
+        may inspect or modify the runner's state in between calls to step().
+        """
+
+        if self._can_launch_more():
+            self._launch_trial()
+        elif self._pending:
+            self._process_events()
+        else:
+            for trial in self._trials:
+                if trial.status == Trial.PENDING:
+                    assert self._has_resources(trial.resources), \
+                        ("Insufficient cluster resources to launch trial",
+                         trial.resources)
+            assert False, "Called step when all trials finished?"
+
+    def get_trials(self):
+        """Returns the list of trials managed by this TrialRunner.
+
+        Note that the caller usually should not mutate trial state directly.
+        """
+
+        return self._trials
+
+    def add_trial(self, trial):
+        """Adds a new trial to this TrialRunner.
+
+        Trials may be added at any time.
+        """
+
+        self._trials.append(trial)
+
+    def debug_string(self):
+        """Returns a human readable message for printing to the console."""
+
+        messages = ["== Status =="]
+        messages.append(
+            "Available: {}".format(self._avail_resources))
+        messages.append(
+            "Committed: {}".format(self._committed_resources))
+        for local_dir in sorted(set([t.local_dir for t in self._trials])):
+            messages.append("Tensorboard logdir: {}".format(local_dir))
+            for t in self._trials:
+                if t.local_dir == local_dir:
+                    messages.append(
+                        " - {}:\t{}".format(t, t.progress_string()))
+        return "\n".join(messages) + "\n"
+
+    def _can_launch_more(self):
+        self._update_avail_resources()
+        trial = self._get_runnable()
+        return trial is not None
+
+    def _launch_trial(self):
+        trial = self._get_runnable()
+        self._commit_resources(trial.resources)
+        try:
+            trial.start()
+            self._pending[trial.train_remote()] = trial
+        except:
+            print("Error starting agent, retrying:", traceback.format_exc())
+            time.sleep(2)
+            trial.stop(error=True)
+            try:
+                trial.start()
+                self._pending[trial.train_remote()] = trial
+            except:
+                print("Error starting agent, abort:", traceback.format_exc())
+                trial.stop(error=True)
+                # note that we don't return the resources, since they may
+                # have been lost
+
+    def _process_events(self):
+        [result_id], _ = ray.wait(self._pending.keys())
+        trial = self._pending[result_id]
+        del self._pending[result_id]
+        try:
+            result = ray.get(result_id)
+            print("result", result)
+            trial.last_result = result
+
+            if trial.should_stop(result):
+                self._return_resources(trial.resources)
+                trial.stop()
+            else:
+                # TODO(rliaw): This implements checkpoint in a blocking manner
+                if trial.should_checkpoint():
+                    trial.checkpoint()
+                self._pending[trial.train_remote()] = trial
+        except:
+            print("Error processing event:", traceback.format_exc())
+            if trial.status == Trial.RUNNING:
+                self._return_resources(trial.resources)
+                trial.stop(error=True)
+
+    def _get_runnable(self):
+        for trial in self._trials:
+            if (trial.status == Trial.PENDING and
+                    self._has_resources(trial.resources)):
+                return trial
+        return None
+
+    def _has_resources(self, resources):
+        cpu_avail = self._avail_resources.cpu - self._committed_resources.cpu
+        gpu_avail = self._avail_resources.gpu - self._committed_resources.gpu
+        assert cpu_avail >= 0 and gpu_avail >= 0
+        return resources.cpu <= cpu_avail and resources.gpu <= gpu_avail
+
+    def _commit_resources(self, resources):
+        self._committed_resources = Resources(
+            self._committed_resources.cpu + resources.cpu,
+            self._committed_resources.gpu + resources.gpu)
+
+    def _return_resources(self, resources):
+        self._committed_resources = Resources(
+            self._committed_resources.cpu - resources.cpu,
+            self._committed_resources.gpu - resources.gpu)
+        assert self._committed_resources.cpu >= 0
+        assert self._committed_resources.gpu >= 0
+
+    def _update_avail_resources(self):
+        clients = ray.global_state.client_table()
+        local_schedulers = [
+            entry for client in clients.values() for entry in client
+            if (entry['ClientType'] == 'local_scheduler' and not
+                entry['Deleted'])
+        ]
+        num_cpus = sum(ls['NumCPUs'] for ls in local_schedulers)
+        num_gpus = sum(ls['NumGPUs'] for ls in local_schedulers)
+        self._avail_resources = Resources(int(num_cpus), int(num_gpus))