From a674ec958c9661b1497e8eb3993dc8e8837ca980 Mon Sep 17 00:00:00 2001 From: Eric Liang Date: Sun, 25 Jun 2017 15:13:03 -0700 Subject: [PATCH] [rllib] Move policy gradient and evolution strategies algorithms from examples/ to ray/rllib/ (#694) * rllib v0 * fix imports * lint * comments * update docs --- doc/source/example-evolution-strategies.rst | 4 +- doc/source/example-policy-gradient.rst | 4 +- examples/policy_gradient/examples/example.py | 194 ------------------ examples/policy_gradient/setup.py | 9 - .../ray/rllib}/__init__.py | 0 python/ray/rllib/common.py | 32 +++ .../rllib/evolution_strategies/__init__.py | 4 + .../evolution_strategies.py | 171 ++++++++------- .../ray/rllib/evolution_strategies/example.py | 40 ++++ .../rllib}/evolution_strategies/optimizers.py | 0 .../rllib}/evolution_strategies/policies.py | 2 +- .../evolution_strategies/tabular_logger.py | 0 .../rllib}/evolution_strategies/tf_util.py | 0 .../ray/rllib}/evolution_strategies/utils.py | 0 .../ray/rllib}/evolution_strategies/viz.py | 0 python/ray/rllib/example.py | 25 +++ python/ray/rllib/policy_gradient/__init__.py | 4 + .../ray/rllib/policy_gradient}/agent.py | 13 +- .../rllib/policy_gradient}/distributions.py | 0 .../ray/rllib/policy_gradient}/env.py | 0 python/ray/rllib/policy_gradient/example.py | 38 ++++ .../ray/rllib/policy_gradient}/filter.py | 0 .../ray/rllib/policy_gradient/loss.py | 4 +- .../rllib/policy_gradient}/models/__init__.py | 0 .../rllib/policy_gradient}/models/fcnet.py | 0 .../policy_gradient}/models/visionnet.py | 0 .../rllib/policy_gradient/policy_gradient.py | 188 +++++++++++++++++ .../ray/rllib/policy_gradient}/rollout.py | 4 +- .../ray/rllib/policy_gradient/test}/test.py | 4 +- .../ray/rllib/policy_gradient}/utils.py | 0 30 files changed, 431 insertions(+), 309 deletions(-) delete mode 100644 examples/policy_gradient/examples/example.py delete mode 100644 examples/policy_gradient/setup.py rename {examples/policy_gradient/reinforce => python/ray/rllib}/__init__.py (100%) create mode 100644 python/ray/rllib/common.py create mode 100644 python/ray/rllib/evolution_strategies/__init__.py rename {examples => python/ray/rllib}/evolution_strategies/evolution_strategies.py (66%) create mode 100755 python/ray/rllib/evolution_strategies/example.py rename {examples => python/ray/rllib}/evolution_strategies/optimizers.py (100%) rename {examples => python/ray/rllib}/evolution_strategies/policies.py (99%) rename {examples => python/ray/rllib}/evolution_strategies/tabular_logger.py (100%) rename {examples => python/ray/rllib}/evolution_strategies/tf_util.py (100%) rename {examples => python/ray/rllib}/evolution_strategies/utils.py (100%) rename {examples => python/ray/rllib}/evolution_strategies/viz.py (100%) create mode 100755 python/ray/rllib/example.py create mode 100644 python/ray/rllib/policy_gradient/__init__.py rename {examples/policy_gradient/reinforce => python/ray/rllib/policy_gradient}/agent.py (96%) rename {examples/policy_gradient/reinforce => python/ray/rllib/policy_gradient}/distributions.py (100%) rename {examples/policy_gradient/reinforce => python/ray/rllib/policy_gradient}/env.py (100%) create mode 100755 python/ray/rllib/policy_gradient/example.py rename {examples/policy_gradient/reinforce => python/ray/rllib/policy_gradient}/filter.py (100%) rename examples/policy_gradient/reinforce/policy.py => python/ray/rllib/policy_gradient/loss.py (93%) rename {examples/policy_gradient/reinforce => python/ray/rllib/policy_gradient}/models/__init__.py (100%) rename {examples/policy_gradient/reinforce => python/ray/rllib/policy_gradient}/models/fcnet.py (100%) rename {examples/policy_gradient/reinforce => python/ray/rllib/policy_gradient}/models/visionnet.py (100%) create mode 100644 python/ray/rllib/policy_gradient/policy_gradient.py rename {examples/policy_gradient/reinforce => python/ray/rllib/policy_gradient}/rollout.py (96%) rename {examples/policy_gradient/tests => python/ray/rllib/policy_gradient/test}/test.py (93%) rename {examples/policy_gradient/reinforce => python/ray/rllib/policy_gradient}/utils.py (100%) diff --git a/doc/source/example-evolution-strategies.rst b/doc/source/example-evolution-strategies.rst index deb45c252..c807500c4 100644 --- a/doc/source/example-evolution-strategies.rst +++ b/doc/source/example-evolution-strategies.rst @@ -11,14 +11,14 @@ To run the application, first install some dependencies. You can view the `code for this example`_. -.. _`code for this example`: https://github.com/ray-project/ray/tree/master/examples/evolution_strategies +.. _`code for this example`: https://github.com/ray-project/ray/tree/master/python/ray/rllib/evolution_strategies The script can be run as follows. Note that the configuration is tuned to work on the ``Humanoid-v1`` gym environment. .. code-block:: bash - python examples/evolution_strategies/evolution_strategies.py + python/ray/rllib/evolution_strategies/example.py At the heart of this example, we define a ``Worker`` class. These workers have a method ``do_rollouts``, which will be used to perform simulate randomly diff --git a/doc/source/example-policy-gradient.rst b/doc/source/example-policy-gradient.rst index a8f1302ec..e47d0ec94 100644 --- a/doc/source/example-policy-gradient.rst +++ b/doc/source/example-policy-gradient.rst @@ -23,7 +23,7 @@ Then you can run the example as follows. .. code-block:: bash - python ray/examples/policy_gradient/examples/example.py --environment=Pong-ram-v3 + python/ray/rllib/policy_gradient/example.py --environment=Pong-ram-v3 This will train an agent on the ``Pong-ram-v3`` Atari environment. You can also try passing in the ``Pong-v0`` environment or the ``CartPole-v0`` environment. @@ -41,4 +41,4 @@ Many of the TensorBoard metrics are also printed to the console, but you might find it easier to visualize and compare between runs using the TensorBoard UI. .. _`TensorFlow with GPU support`: https://www.tensorflow.org/install/ -.. _`code for this example`: https://github.com/ray-project/ray/tree/master/examples/policy_gradient +.. _`code for this example`: https://github.com/ray-project/ray/tree/master/python/ray/rllib/policy_gradient diff --git a/examples/policy_gradient/examples/example.py b/examples/policy_gradient/examples/example.py deleted file mode 100644 index 0e65ec887..000000000 --- a/examples/policy_gradient/examples/example.py +++ /dev/null @@ -1,194 +0,0 @@ -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -from datetime import datetime - -import argparse -import time - -import ray -import numpy as np -import tensorflow as tf - -from reinforce.env import (NoPreprocessor, AtariRamPreprocessor, - AtariPixelPreprocessor) -from reinforce.agent import Agent, RemoteAgent -from reinforce.rollout import collect_samples -from reinforce.utils import shuffle - - -config = {"kl_coeff": 0.2, - "num_sgd_iter": 30, - "max_iterations": 1000, - "sgd_stepsize": 5e-5, - # TODO(pcm): Expose the choice between gpus and cpus - # as a command line argument. - "devices": ["/cpu:%d" % i for i in range(4)], - "tf_session_args": { - "device_count": {"CPU": 4}, - "log_device_placement": False, - "allow_soft_placement": True, - }, - "sgd_batchsize": 128, # total size across all devices - "entropy_coeff": 0.0, - "clip_param": 0.3, - "kl_target": 0.01, - "timesteps_per_batch": 40000, - "num_agents": 5, - "tensorboard_log_dir": "/tmp/ray", - "full_trace_nth_sgd_batch": -1, - "full_trace_data_load": False, - "model_checkpoint_file": "/tmp/iteration-%s.ckpt"} - - -if __name__ == "__main__": - parser = argparse.ArgumentParser(description="Run the policy gradient " - "algorithm.") - parser.add_argument("--environment", default="Pong-v0", type=str, - help="The gym environment to use.") - parser.add_argument("--redis-address", default=None, type=str, - help="The Redis address of the cluster.") - parser.add_argument("--use-tf-debugger", default=False, type=bool, - help="Run the script inside of tf-dbg.") - parser.add_argument("--load-checkpoint", default=None, type=str, - help="Continue training from a checkpoint.") - - args = parser.parse_args() - config["use_tf_debugger"] = args.use_tf_debugger - - ray.init(redis_address=args.redis_address) - - mdp_name = args.environment - if args.environment == "Pong-v0": - preprocessor = AtariPixelPreprocessor() - elif mdp_name == "Pong-ram-v3": - preprocessor = AtariRamPreprocessor() - elif mdp_name == "CartPole-v0": - preprocessor = NoPreprocessor() - elif mdp_name == "Walker2d-v1": - preprocessor = NoPreprocessor() - else: - print("No environment was chosen, so defaulting to Pong-v0.") - mdp_name = "Pong-v0" - preprocessor = AtariPixelPreprocessor() - - print("Using the environment {}.".format(mdp_name)) - agents = [RemoteAgent.remote(mdp_name, 1, preprocessor, config, True) - for _ in range(config["num_agents"])] - agent = Agent(mdp_name, 1, preprocessor, config, False) - - kl_coeff = config["kl_coeff"] - - file_writer = tf.summary.FileWriter( - "{}/trpo_{}_{}".format( - config["tensorboard_log_dir"], mdp_name, - str(datetime.today()).replace(" ", "_")), - agent.sess.graph) - - global_step = 0 - - saver = tf.train.Saver(max_to_keep=None) - if args.load_checkpoint: - saver.restore(agent.sess, args.load_checkpoint) - - for j in range(config["max_iterations"]): - iter_start = time.time() - print("\n== iteration", j) - if config["model_checkpoint_file"]: - checkpoint_path = saver.save( - agent.sess, config["model_checkpoint_file"] % j) - print("Checkpoint saved in file: %s" % checkpoint_path) - checkpointing_end = time.time() - weights = ray.put(agent.get_weights()) - [a.load_weights.remote(weights) for a in agents] - trajectory, total_reward, traj_len_mean = collect_samples( - agents, config["timesteps_per_batch"], 0.995, 1.0, 2000) - print("total reward is ", total_reward) - print("trajectory length mean is ", traj_len_mean) - print("timesteps:", trajectory["dones"].shape[0]) - traj_stats = tf.Summary(value=[ - tf.Summary.Value( - tag="policy_gradient/rollouts/mean_reward", - simple_value=total_reward), - tf.Summary.Value( - tag="policy_gradient/rollouts/traj_len_mean", - simple_value=traj_len_mean)]) - file_writer.add_summary(traj_stats, global_step) - global_step += 1 - trajectory["advantages"] = ((trajectory["advantages"] - - trajectory["advantages"].mean()) / - trajectory["advantages"].std()) - rollouts_end = time.time() - print("Computing policy (iterations=" + str(config["num_sgd_iter"]) + - ", stepsize=" + str(config["sgd_stepsize"]) + "):") - names = ["iter", "loss", "kl", "entropy"] - print(("{:>15}" * len(names)).format(*names)) - num_devices = len(config["devices"]) - trajectory = shuffle(trajectory) - shuffle_end = time.time() - tuples_per_device = agent.load_data( - trajectory, j == 0 and config["full_trace_data_load"]) - load_end = time.time() - checkpointing_time = checkpointing_end - iter_start - rollouts_time = rollouts_end - checkpointing_end - shuffle_time = shuffle_end - rollouts_end - load_time = load_end - shuffle_end - sgd_time = 0 - for i in range(config["num_sgd_iter"]): - sgd_start = time.time() - batch_index = 0 - num_batches = int(tuples_per_device) // int(agent.per_device_batch_size) - loss, kl, entropy = [], [], [] - permutation = np.random.permutation(num_batches) - while batch_index < num_batches: - full_trace = ( - i == 0 and j == 0 and - batch_index == config["full_trace_nth_sgd_batch"]) - batch_loss, batch_kl, batch_entropy = agent.run_sgd_minibatch( - permutation[batch_index] * agent.per_device_batch_size, - kl_coeff, full_trace, file_writer) - loss.append(batch_loss) - kl.append(batch_kl) - entropy.append(batch_entropy) - batch_index += 1 - loss = np.mean(loss) - kl = np.mean(kl) - entropy = np.mean(entropy) - sgd_end = time.time() - print("{:>15}{:15.5e}{:15.5e}{:15.5e}".format(i, loss, kl, entropy)) - - values = [] - if i == config["num_sgd_iter"] - 1: - metric_prefix = "policy_gradient/sgd/final_iter/" - values.append(tf.Summary.Value( - tag=metric_prefix + "kl_coeff", - simple_value=kl_coeff)) - else: - metric_prefix = "policy_gradient/sgd/intermediate_iters/" - values.extend([ - tf.Summary.Value( - tag=metric_prefix + "mean_entropy", - simple_value=entropy), - tf.Summary.Value( - tag=metric_prefix + "mean_loss", - simple_value=loss), - tf.Summary.Value( - tag=metric_prefix + "mean_kl", - simple_value=kl)]) - sgd_stats = tf.Summary(value=values) - file_writer.add_summary(sgd_stats, global_step) - global_step += 1 - sgd_time += sgd_end - sgd_start - if kl > 2.0 * config["kl_target"]: - kl_coeff *= 1.5 - elif kl < 0.5 * config["kl_target"]: - kl_coeff *= 0.5 - print("kl div:", kl) - print("kl coeff:", kl_coeff) - print("checkpointing time:", checkpointing_time) - print("rollouts time:", rollouts_time) - print("shuffle time:", shuffle_time) - print("load time:", load_time) - print("sgd time:", sgd_time) - print("sgd examples/s:", len(trajectory["observations"]) / sgd_time) diff --git a/examples/policy_gradient/setup.py b/examples/policy_gradient/setup.py deleted file mode 100644 index bc02ddb21..000000000 --- a/examples/policy_gradient/setup.py +++ /dev/null @@ -1,9 +0,0 @@ -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -from setuptools import setup, find_packages - -setup(name="reinforce", - version="0.0.1", - packages=find_packages()) diff --git a/examples/policy_gradient/reinforce/__init__.py b/python/ray/rllib/__init__.py similarity index 100% rename from examples/policy_gradient/reinforce/__init__.py rename to python/ray/rllib/__init__.py diff --git a/python/ray/rllib/common.py b/python/ray/rllib/common.py new file mode 100644 index 000000000..050ccf535 --- /dev/null +++ b/python/ray/rllib/common.py @@ -0,0 +1,32 @@ +from collections import namedtuple + + +TrainingResult = namedtuple("TrainingResult", [ + "training_iteration", + "episode_reward_mean", + "episode_len_mean", +]) + + +class Algorithm(object): + """All RLlib algorithms extend this base class. + + Algorithm objects retain internal model state between calls to train(), so + you should create a new algorithm instance for each training session. + + TODO(ekl): support checkpoint / restore of training state. + """ + + def __init__(self, env_name, config): + self.env_name = env_name + self.config = config + + def train(self): + """ + Runs one logical iteration of training. + + Returns: + A TrainingResult that describes training progress. + """ + + raise NotImplementedError diff --git a/python/ray/rllib/evolution_strategies/__init__.py b/python/ray/rllib/evolution_strategies/__init__.py new file mode 100644 index 000000000..6b064796e --- /dev/null +++ b/python/ray/rllib/evolution_strategies/__init__.py @@ -0,0 +1,4 @@ +from ray.rllib.evolution_strategies.evolution_strategies import ( + EvolutionStrategies, DEFAULT_CONFIG) + +__all__ = ["EvolutionStrategies", "DEFAULT_CONFIG"] diff --git a/examples/evolution_strategies/evolution_strategies.py b/python/ray/rllib/evolution_strategies/evolution_strategies.py similarity index 66% rename from examples/evolution_strategies/evolution_strategies.py rename to python/ray/rllib/evolution_strategies/evolution_strategies.py index 3d5296988..b9c6bbe42 100644 --- a/examples/evolution_strategies/evolution_strategies.py +++ b/python/ray/rllib/evolution_strategies/evolution_strategies.py @@ -5,25 +5,26 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -import argparse from collections import namedtuple import gym import numpy as np import os -import ray import time -import optimizers -import policies -import tabular_logger as tlogger -import tf_util -import utils +import ray +from ray.rllib.common import Algorithm, TrainingResult + +from ray.rllib.evolution_strategies import optimizers +from ray.rllib.evolution_strategies import policies +from ray.rllib.evolution_strategies import tabular_logger as tlogger +from ray.rllib.evolution_strategies import tf_util +from ray.rllib.evolution_strategies import utils Config = namedtuple("Config", [ "l2coeff", "noise_stdev", "episodes_per_batch", "timesteps_per_batch", "calc_obstat_prob", "eval_prob", "snapshot_freq", "return_proc_mode", - "episode_cutoff_mode" + "episode_cutoff_mode", "num_workers", "stepsize" ]) Result = namedtuple("Result", [ @@ -32,6 +33,20 @@ Result = namedtuple("Result", [ ]) +DEFAULT_CONFIG = Config( + l2coeff=0.005, + noise_stdev=0.02, + episodes_per_batch=10000, + timesteps_per_batch=100000, + calc_obstat_prob=0.01, + eval_prob=0, + snapshot_freq=0, + return_proc_mode="centered_rank", + episode_cutoff_mode="env_default", + num_workers=10, + stepsize=.01) + + @ray.remote def create_shared_noise(): """Create a large array of noise to be shared by all workers.""" @@ -135,72 +150,46 @@ class Worker(object): ob_count=task_ob_stat.count) -if __name__ == "__main__": - parser = argparse.ArgumentParser(description="Train an RL agent on Pong.") - parser.add_argument("--num-workers", default=10, type=int, - help=("The number of actors to create in aggregate " - "across the cluster.")) - parser.add_argument("--env-name", default="Pendulum-v0", type=str, - help="The name of the gym environment to use.") - parser.add_argument("--stepsize", default=0.01, type=float, - help="The stepsize to use.") - parser.add_argument("--redis-address", default=None, type=str, - help="The Redis address of the cluster.") +class EvolutionStrategies(Algorithm): + def __init__(self, env_name, config): + Algorithm.__init__(self, env_name, config) - args = parser.parse_args() - num_workers = args.num_workers - env_name = args.env_name - stepsize = args.stepsize + policy_params = { + "ac_bins": "continuous:", + "ac_noise_std": 0.01, + "nonlin_type": "tanh", + "hidden_dims": [256, 256], + "connection_type": "ff" + } - ray.init(redis_address=args.redis_address, - num_workers=(0 if args.redis_address is None else None)) + # Create the shared noise table. + print("Creating shared noise table.") + noise_id = create_shared_noise.remote() + self.noise = SharedNoiseTable(ray.get(noise_id)) - config = Config(l2coeff=0.005, - noise_stdev=0.02, - episodes_per_batch=10000, - timesteps_per_batch=100000, - calc_obstat_prob=0.01, - eval_prob=0, - snapshot_freq=20, - return_proc_mode="centered_rank", - episode_cutoff_mode="env_default") + # Create the actors. + print("Creating actors.") + self.workers = [Worker.remote(config, policy_params, env_name, noise_id) + for _ in range(config.num_workers)] - policy_params = { - "ac_bins": "continuous:", - "ac_noise_std": 0.01, - "nonlin_type": "tanh", - "hidden_dims": [256, 256], - "connection_type": "ff" - } + env = gym.make(env_name) + utils.make_session(single_threaded=False) + self.policy = policies.MujocoPolicy( + env.observation_space, env.action_space, **policy_params) + tf_util.initialize() + self.optimizer = optimizers.Adam(self.policy, config.stepsize) + self.ob_stat = utils.RunningStat(env.observation_space.shape, eps=1e-2) - # Create the shared noise table. - print("Creating shared noise table.") - noise_id = create_shared_noise.remote() - noise = SharedNoiseTable(ray.get(noise_id)) + self.episodes_so_far = 0 + self.timesteps_so_far = 0 + self.tstart = time.time() + self.iteration = 0 - # Create the actors. - print("Creating actors.") - workers = [Worker.remote(config, policy_params, env_name, noise_id) - for _ in range(num_workers)] + def train(self): + config = self.config - env = gym.make(env_name) - sess = utils.make_session(single_threaded=False) - policy = policies.MujocoPolicy(env.observation_space, env.action_space, - **policy_params) - tf_util.initialize() - optimizer = optimizers.Adam(policy, stepsize) - - ob_stat = utils.RunningStat(env.observation_space.shape, eps=1e-2) - - episodes_so_far = 0 - timesteps_so_far = 0 - tstart = time.time() - - iteration = 0 - - while True: step_tstart = time.time() - theta = policy.get_trainable_flat() + theta = self.policy.get_trainable_flat() assert theta.dtype == np.float32 # Put the current policy weights in the object store. @@ -209,8 +198,9 @@ if __name__ == "__main__": # weights. rollout_ids = [worker.do_rollouts.remote( theta_id, - ob_stat.mean if policy.needs_ob_stat else None, - ob_stat.std if policy.needs_ob_stat else None) for worker in workers] + self.ob_stat.mean if self.policy.needs_ob_stat else None, + self.ob_stat.std if self.policy.needs_ob_stat else None) + for worker in self.workers] # Get the results of the rollouts. results = ray.get(rollout_ids) @@ -227,13 +217,13 @@ if __name__ == "__main__": result_num_eps = result.lengths_n2.size result_num_timesteps = result.lengths_n2.sum() - episodes_so_far += result_num_eps - timesteps_so_far += result_num_timesteps + self.episodes_so_far += result_num_eps + self.timesteps_so_far += result_num_timesteps curr_task_results.append(result) # Update ob stats. - if policy.needs_ob_stat and result.ob_count > 0: - ob_stat.increment(result.ob_sum, result.ob_sumsq, result.ob_count) + if self.policy.needs_ob_stat and result.ob_count > 0: + self.ob_stat.increment(result.ob_sum, result.ob_sumsq, result.ob_count) ob_count_this_batch += result.ob_count # Assemble the results. @@ -251,44 +241,47 @@ if __name__ == "__main__": # Compute and take a step. g, count = utils.batched_weighted_sum( proc_returns_n2[:, 0] - proc_returns_n2[:, 1], - (noise.get(idx, policy.num_params) for idx in noise_inds_n), + (self.noise.get(idx, self.policy.num_params) for idx in noise_inds_n), batch_size=500) g /= returns_n2.size - assert (g.shape == (policy.num_params,) and g.dtype == np.float32 and + assert (g.shape == (self.policy.num_params,) and g.dtype == np.float32 and count == len(noise_inds_n)) - update_ratio = optimizer.update(-g + config.l2coeff * theta) + update_ratio = self.optimizer.update(-g + config.l2coeff * theta) # Update ob stat (we're never running the policy in the master, but we # might be snapshotting the policy). - if policy.needs_ob_stat: - policy.set_ob_stat(ob_stat.mean, ob_stat.std) + if self.policy.needs_ob_stat: + self.policy.set_ob_stat(self.ob_stat.mean, self.ob_stat.std) step_tend = time.time() tlogger.record_tabular("EpRewMean", returns_n2.mean()) tlogger.record_tabular("EpRewStd", returns_n2.std()) tlogger.record_tabular("EpLenMean", lengths_n2.mean()) - tlogger.record_tabular("Norm", - float(np.square(policy.get_trainable_flat()).sum())) + tlogger.record_tabular( + "Norm", float(np.square(self.policy.get_trainable_flat()).sum())) tlogger.record_tabular("GradNorm", float(np.square(g).sum())) tlogger.record_tabular("UpdateRatio", float(update_ratio)) tlogger.record_tabular("EpisodesThisIter", lengths_n2.size) - tlogger.record_tabular("EpisodesSoFar", episodes_so_far) + tlogger.record_tabular("EpisodesSoFar", self.episodes_so_far) tlogger.record_tabular("TimestepsThisIter", lengths_n2.sum()) - tlogger.record_tabular("TimestepsSoFar", timesteps_so_far) + tlogger.record_tabular("TimestepsSoFar", self.timesteps_so_far) tlogger.record_tabular("ObCount", ob_count_this_batch) tlogger.record_tabular("TimeElapsedThisIter", step_tend - step_tstart) - tlogger.record_tabular("TimeElapsed", step_tend - tstart) + tlogger.record_tabular("TimeElapsed", step_tend - self.tstart) tlogger.dump_tabular() - if config.snapshot_freq != 0 and iteration % config.snapshot_freq == 0: - filename = os.path.join("/tmp", - "snapshot_iter{:05d}.h5".format(iteration)) - assert not os.path.exists(filename) - policy.save(filename) - tlogger.log("Saved snapshot {}".format(filename)) + if (config.snapshot_freq != 0 and + self.iteration % config.snapshot_freq == 0): + filename = os.path.join( + "/tmp", "snapshot_iter{:05d}.h5".format(self.iteration)) + assert not os.path.exists(filename) + self.policy.save(filename) + tlogger.log("Saved snapshot {}".format(filename)) - iteration += 1 + res = TrainingResult(self.iteration, returns_n2.mean(), lengths_n2.mean()) + self.iteration += 1 + return res diff --git a/python/ray/rllib/evolution_strategies/example.py b/python/ray/rllib/evolution_strategies/example.py new file mode 100755 index 000000000..5af53e367 --- /dev/null +++ b/python/ray/rllib/evolution_strategies/example.py @@ -0,0 +1,40 @@ +#!/usr/bin/env python + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import argparse + +import ray +from ray.rllib.evolution_strategies import EvolutionStrategies, DEFAULT_CONFIG + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Train an RL agent on Pong.") + parser.add_argument("--num-workers", default=10, type=int, + help=("The number of actors to create in aggregate " + "across the cluster.")) + parser.add_argument("--env-name", default="Pendulum-v0", type=str, + help="The name of the gym environment to use.") + parser.add_argument("--stepsize", default=0.01, type=float, + help="The stepsize to use.") + parser.add_argument("--redis-address", default=None, type=str, + help="The Redis address of the cluster.") + + args = parser.parse_args() + num_workers = args.num_workers + env_name = args.env_name + stepsize = args.stepsize + + ray.init(redis_address=args.redis_address, + num_workers=(0 if args.redis_address is None else None)) + + config = DEFAULT_CONFIG._replace( + num_workers=num_workers, + stepsize=stepsize) + + alg = EvolutionStrategies(env_name, config) + while True: + result = alg.train() + print("current status: {}".format(result)) diff --git a/examples/evolution_strategies/optimizers.py b/python/ray/rllib/evolution_strategies/optimizers.py similarity index 100% rename from examples/evolution_strategies/optimizers.py rename to python/ray/rllib/evolution_strategies/optimizers.py diff --git a/examples/evolution_strategies/policies.py b/python/ray/rllib/evolution_strategies/policies.py similarity index 99% rename from examples/evolution_strategies/policies.py rename to python/ray/rllib/evolution_strategies/policies.py index 8f9a2ded2..7997fb8ce 100644 --- a/examples/evolution_strategies/policies.py +++ b/python/ray/rllib/evolution_strategies/policies.py @@ -12,7 +12,7 @@ import h5py import numpy as np import tensorflow as tf -import tf_util as U +from ray.rllib.evolution_strategies import tf_util as U logger = logging.getLogger(__name__) diff --git a/examples/evolution_strategies/tabular_logger.py b/python/ray/rllib/evolution_strategies/tabular_logger.py similarity index 100% rename from examples/evolution_strategies/tabular_logger.py rename to python/ray/rllib/evolution_strategies/tabular_logger.py diff --git a/examples/evolution_strategies/tf_util.py b/python/ray/rllib/evolution_strategies/tf_util.py similarity index 100% rename from examples/evolution_strategies/tf_util.py rename to python/ray/rllib/evolution_strategies/tf_util.py diff --git a/examples/evolution_strategies/utils.py b/python/ray/rllib/evolution_strategies/utils.py similarity index 100% rename from examples/evolution_strategies/utils.py rename to python/ray/rllib/evolution_strategies/utils.py diff --git a/examples/evolution_strategies/viz.py b/python/ray/rllib/evolution_strategies/viz.py similarity index 100% rename from examples/evolution_strategies/viz.py rename to python/ray/rllib/evolution_strategies/viz.py diff --git a/python/ray/rllib/example.py b/python/ray/rllib/example.py new file mode 100755 index 000000000..8d2b353a7 --- /dev/null +++ b/python/ray/rllib/example.py @@ -0,0 +1,25 @@ +#!/usr/bin/env python +"""Demonstrates the RLlib algorithm API through a simple bakeoff.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import ray +import ray.rllib.evolution_strategies as es +import ray.rllib.policy_gradient as pg + + +if __name__ == "__main__": + ray.init() + + # TODO(ekl): get the algorithms working on a common set of envs + env_name = "CartPole-v0" + alg1 = es.EvolutionStrategies(env_name, es.DEFAULT_CONFIG) + alg2 = pg.PolicyGradient(env_name, pg.DEFAULT_CONFIG) + + while True: + r1 = alg1.train() + r2 = alg2.train() + print("evolution strategies: {}".format(r1)) + print("policy gradient: {}".format(r2)) diff --git a/python/ray/rllib/policy_gradient/__init__.py b/python/ray/rllib/policy_gradient/__init__.py new file mode 100644 index 000000000..bcd558423 --- /dev/null +++ b/python/ray/rllib/policy_gradient/__init__.py @@ -0,0 +1,4 @@ +from ray.rllib.policy_gradient.policy_gradient import ( + PolicyGradient, DEFAULT_CONFIG) + +__all__ = ["PolicyGradient", "DEFAULT_CONFIG"] diff --git a/examples/policy_gradient/reinforce/agent.py b/python/ray/rllib/policy_gradient/agent.py similarity index 96% rename from examples/policy_gradient/reinforce/agent.py rename to python/ray/rllib/policy_gradient/agent.py index 380519742..cfe1386e1 100644 --- a/examples/policy_gradient/reinforce/agent.py +++ b/python/ray/rllib/policy_gradient/agent.py @@ -13,12 +13,13 @@ from tensorflow.python import debug as tf_debug import ray -from reinforce.distributions import Categorical, DiagGaussian -from reinforce.env import BatchedEnv -from reinforce.policy import ProximalPolicyLoss -from reinforce.filter import MeanStdFilter -from reinforce.rollout import rollouts, add_advantage_values -from reinforce.utils import make_divisible_by, average_gradients +from ray.rllib.policy_gradient.distributions import Categorical, DiagGaussian +from ray.rllib.policy_gradient.env import BatchedEnv +from ray.rllib.policy_gradient.loss import ProximalPolicyLoss +from ray.rllib.policy_gradient.filter import MeanStdFilter +from ray.rllib.policy_gradient.rollout import rollouts, add_advantage_values +from ray.rllib.policy_gradient.utils import ( + make_divisible_by, average_gradients) # TODO(pcm): Make sure that both observation_filter and reward_filter # are correctly handled, i.e. (a) the values are accumulated accross diff --git a/examples/policy_gradient/reinforce/distributions.py b/python/ray/rllib/policy_gradient/distributions.py similarity index 100% rename from examples/policy_gradient/reinforce/distributions.py rename to python/ray/rllib/policy_gradient/distributions.py diff --git a/examples/policy_gradient/reinforce/env.py b/python/ray/rllib/policy_gradient/env.py similarity index 100% rename from examples/policy_gradient/reinforce/env.py rename to python/ray/rllib/policy_gradient/env.py diff --git a/python/ray/rllib/policy_gradient/example.py b/python/ray/rllib/policy_gradient/example.py new file mode 100755 index 000000000..dc81ced01 --- /dev/null +++ b/python/ray/rllib/policy_gradient/example.py @@ -0,0 +1,38 @@ +#!/usr/bin/env python + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import argparse + +import ray +from ray.rllib.policy_gradient import PolicyGradient, DEFAULT_CONFIG + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Run the policy gradient " + "algorithm.") + parser.add_argument("--environment", default="Pong-v0", type=str, + help="The gym environment to use.") + parser.add_argument("--redis-address", default=None, type=str, + help="The Redis address of the cluster.") + parser.add_argument("--use-tf-debugger", default=False, type=bool, + help="Run the script inside of tf-dbg.") + parser.add_argument("--load-checkpoint", default=None, type=str, + help="Continue training from a checkpoint.") + + args = parser.parse_args() + config = DEFAULT_CONFIG.copy() + config["use_tf_debugger"] = args.use_tf_debugger + if args.load_checkpoint: + config["load_checkpoint"] = args.load_checkpoint + + ray.init(redis_address=args.redis_address) + + alg = PolicyGradient(args.environment, config) + result = alg.train() + while result.training_iteration < config["max_iterations"]: + print("\n== iteration", result.training_iteration) + result = alg.train() + print("current status: {}".format(result)) diff --git a/examples/policy_gradient/reinforce/filter.py b/python/ray/rllib/policy_gradient/filter.py similarity index 100% rename from examples/policy_gradient/reinforce/filter.py rename to python/ray/rllib/policy_gradient/filter.py diff --git a/examples/policy_gradient/reinforce/policy.py b/python/ray/rllib/policy_gradient/loss.py similarity index 93% rename from examples/policy_gradient/reinforce/policy.py rename to python/ray/rllib/policy_gradient/loss.py index db378113d..2f2129150 100644 --- a/examples/policy_gradient/reinforce/policy.py +++ b/python/ray/rllib/policy_gradient/loss.py @@ -4,8 +4,8 @@ from __future__ import print_function import gym.spaces import tensorflow as tf -from reinforce.models.visionnet import vision_net -from reinforce.models.fcnet import fc_net +from ray.rllib.policy_gradient.models.visionnet import vision_net +from ray.rllib.policy_gradient.models.fcnet import fc_net class ProximalPolicyLoss(object): diff --git a/examples/policy_gradient/reinforce/models/__init__.py b/python/ray/rllib/policy_gradient/models/__init__.py similarity index 100% rename from examples/policy_gradient/reinforce/models/__init__.py rename to python/ray/rllib/policy_gradient/models/__init__.py diff --git a/examples/policy_gradient/reinforce/models/fcnet.py b/python/ray/rllib/policy_gradient/models/fcnet.py similarity index 100% rename from examples/policy_gradient/reinforce/models/fcnet.py rename to python/ray/rllib/policy_gradient/models/fcnet.py diff --git a/examples/policy_gradient/reinforce/models/visionnet.py b/python/ray/rllib/policy_gradient/models/visionnet.py similarity index 100% rename from examples/policy_gradient/reinforce/models/visionnet.py rename to python/ray/rllib/policy_gradient/models/visionnet.py diff --git a/python/ray/rllib/policy_gradient/policy_gradient.py b/python/ray/rllib/policy_gradient/policy_gradient.py new file mode 100644 index 000000000..7850af337 --- /dev/null +++ b/python/ray/rllib/policy_gradient/policy_gradient.py @@ -0,0 +1,188 @@ +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from datetime import datetime +import time + +import numpy as np +import tensorflow as tf + +import ray +from ray.rllib.common import Algorithm, TrainingResult +from ray.rllib.policy_gradient.agent import Agent, RemoteAgent +from ray.rllib.policy_gradient.env import ( + NoPreprocessor, AtariRamPreprocessor, AtariPixelPreprocessor) +from ray.rllib.policy_gradient.rollout import collect_samples +from ray.rllib.policy_gradient.utils import shuffle + + +DEFAULT_CONFIG = { + "kl_coeff": 0.2, + "num_sgd_iter": 30, + "max_iterations": 1000, + "sgd_stepsize": 5e-5, + # TODO(pcm): Expose the choice between gpus and cpus + # as a command line argument. + "devices": ["/cpu:%d" % i for i in range(4)], + "tf_session_args": { + "device_count": {"CPU": 4}, + "log_device_placement": False, + "allow_soft_placement": True, + }, + "sgd_batchsize": 128, # total size across all devices + "entropy_coeff": 0.0, + "clip_param": 0.3, + "kl_target": 0.01, + "timesteps_per_batch": 40000, + "num_agents": 5, + "tensorboard_log_dir": "/tmp/ray", + "full_trace_nth_sgd_batch": -1, + "full_trace_data_load": False, + "use_tf_debugger": False, + "model_checkpoint_file": "/tmp/iteration-%s.ckpt"} + + +class PolicyGradient(Algorithm): + def __init__(self, env_name, config): + Algorithm.__init__(self, env_name, config) + + # TODO(ekl) the preprocessor should be associated with the env elsewhere + if self.env_name == "Pong-v0": + preprocessor = AtariPixelPreprocessor() + elif self.env_name == "Pong-ram-v3": + preprocessor = AtariRamPreprocessor() + elif self.env_name == "CartPole-v0": + preprocessor = NoPreprocessor() + elif self.env_name == "Walker2d-v1": + preprocessor = NoPreprocessor() + else: + preprocessor = AtariPixelPreprocessor() + + self.preprocessor = preprocessor + self.global_step = 0 + self.j = 0 + self.kl_coeff = config["kl_coeff"] + self.model = Agent( + self.env_name, 1, self.preprocessor, self.config, False) + self.agents = [ + RemoteAgent.remote( + self.env_name, 1, self.preprocessor, self.config, True) + for _ in range(config["num_agents"])] + + def train(self): + agents = self.agents + config = self.config + model = self.model + j = self.j + self.j += 1 + + saver = tf.train.Saver(max_to_keep=None) + if "load_checkpoint" in config: + saver.restore(model.sess, config["load_checkpoint"]) + + file_writer = tf.summary.FileWriter( + "{}/trpo_{}_{}".format( + config["tensorboard_log_dir"], self.env_name, + str(datetime.today()).replace(" ", "_")), + model.sess.graph) + iter_start = time.time() + if config["model_checkpoint_file"]: + checkpoint_path = saver.save( + model.sess, config["model_checkpoint_file"] % j) + print("Checkpoint saved in file: %s" % checkpoint_path) + checkpointing_end = time.time() + weights = ray.put(model.get_weights()) + [a.load_weights.remote(weights) for a in agents] + trajectory, total_reward, traj_len_mean = collect_samples( + agents, config["timesteps_per_batch"], 0.995, 1.0, 2000) + print("total reward is ", total_reward) + print("trajectory length mean is ", traj_len_mean) + print("timesteps:", trajectory["dones"].shape[0]) + traj_stats = tf.Summary(value=[ + tf.Summary.Value( + tag="policy_gradient/rollouts/mean_reward", + simple_value=total_reward), + tf.Summary.Value( + tag="policy_gradient/rollouts/traj_len_mean", + simple_value=traj_len_mean)]) + file_writer.add_summary(traj_stats, self.global_step) + self.global_step += 1 + trajectory["advantages"] = ((trajectory["advantages"] - + trajectory["advantages"].mean()) / + trajectory["advantages"].std()) + rollouts_end = time.time() + print("Computing policy (iterations=" + str(config["num_sgd_iter"]) + + ", stepsize=" + str(config["sgd_stepsize"]) + "):") + names = ["iter", "loss", "kl", "entropy"] + print(("{:>15}" * len(names)).format(*names)) + trajectory = shuffle(trajectory) + shuffle_end = time.time() + tuples_per_device = model.load_data( + trajectory, j == 0 and config["full_trace_data_load"]) + load_end = time.time() + checkpointing_time = checkpointing_end - iter_start + rollouts_time = rollouts_end - checkpointing_end + shuffle_time = shuffle_end - rollouts_end + load_time = load_end - shuffle_end + sgd_time = 0 + for i in range(config["num_sgd_iter"]): + sgd_start = time.time() + batch_index = 0 + num_batches = int(tuples_per_device) // int(model.per_device_batch_size) + loss, kl, entropy = [], [], [] + permutation = np.random.permutation(num_batches) + while batch_index < num_batches: + full_trace = ( + i == 0 and j == 0 and + batch_index == config["full_trace_nth_sgd_batch"]) + batch_loss, batch_kl, batch_entropy = model.run_sgd_minibatch( + permutation[batch_index] * model.per_device_batch_size, + self.kl_coeff, full_trace, file_writer) + loss.append(batch_loss) + kl.append(batch_kl) + entropy.append(batch_entropy) + batch_index += 1 + loss = np.mean(loss) + kl = np.mean(kl) + entropy = np.mean(entropy) + sgd_end = time.time() + print("{:>15}{:15.5e}{:15.5e}{:15.5e}".format(i, loss, kl, entropy)) + + values = [] + if i == config["num_sgd_iter"] - 1: + metric_prefix = "policy_gradient/sgd/final_iter/" + values.append(tf.Summary.Value( + tag=metric_prefix + "kl_coeff", + simple_value=self.kl_coeff)) + else: + metric_prefix = "policy_gradient/sgd/intermediate_iters/" + values.extend([ + tf.Summary.Value( + tag=metric_prefix + "mean_entropy", + simple_value=entropy), + tf.Summary.Value( + tag=metric_prefix + "mean_loss", + simple_value=loss), + tf.Summary.Value( + tag=metric_prefix + "mean_kl", + simple_value=kl)]) + sgd_stats = tf.Summary(value=values) + file_writer.add_summary(sgd_stats, self.global_step) + self.global_step += 1 + sgd_time += sgd_end - sgd_start + if kl > 2.0 * config["kl_target"]: + self.kl_coeff *= 1.5 + elif kl < 0.5 * config["kl_target"]: + self.kl_coeff *= 0.5 + + print("kl div:", kl) + print("kl coeff:", self.kl_coeff) + print("checkpointing time:", checkpointing_time) + print("rollouts time:", rollouts_time) + print("shuffle time:", shuffle_time) + print("load time:", load_time) + print("sgd time:", sgd_time) + print("sgd examples/s:", len(trajectory["observations"]) / sgd_time) + + return TrainingResult(j, total_reward, traj_len_mean) diff --git a/examples/policy_gradient/reinforce/rollout.py b/python/ray/rllib/policy_gradient/rollout.py similarity index 96% rename from examples/policy_gradient/reinforce/rollout.py rename to python/ray/rllib/policy_gradient/rollout.py index 092b45571..aed531ec2 100644 --- a/examples/policy_gradient/reinforce/rollout.py +++ b/python/ray/rllib/policy_gradient/rollout.py @@ -5,8 +5,8 @@ from __future__ import print_function import numpy as np import ray -from reinforce.filter import NoFilter -from reinforce.utils import flatten, concatenate +from ray.rllib.policy_gradient.filter import NoFilter +from ray.rllib.policy_gradient.utils import flatten, concatenate def rollouts(policy, env, horizon, observation_filter=NoFilter(), diff --git a/examples/policy_gradient/tests/test.py b/python/ray/rllib/policy_gradient/test/test.py similarity index 93% rename from examples/policy_gradient/tests/test.py rename to python/ray/rllib/policy_gradient/test/test.py index c1db6ddc5..0743bdc5d 100644 --- a/examples/policy_gradient/tests/test.py +++ b/python/ray/rllib/policy_gradient/test/test.py @@ -7,8 +7,8 @@ import numpy as np import tensorflow as tf from numpy.testing import assert_allclose -from reinforce.distributions import Categorical -from reinforce.utils import flatten, concatenate +from ray.rllib.policy_gradient.distributions import Categorical +from ray.rllib.policy_gradient.utils import flatten, concatenate class DistibutionsTest(unittest.TestCase): diff --git a/examples/policy_gradient/reinforce/utils.py b/python/ray/rllib/policy_gradient/utils.py similarity index 100% rename from examples/policy_gradient/reinforce/utils.py rename to python/ray/rllib/policy_gradient/utils.py