diff --git a/doc/source/example-a3c.rst b/doc/source/example-a3c.rst new file mode 100644 index 000000000..2f249a656 --- /dev/null +++ b/doc/source/example-a3c.rst @@ -0,0 +1,146 @@ +Asynchronous Advantage Actor Critic (A3C) +========================================= + +This document walks through `A3C`_, a state-of-the-art reinforcement learning +algorithm. In this example, we adapt the OpenAI `Universe Starter Agent`_ +implementation of A3C to use Ray. + +View the `code for this example`_. + +.. _`A3C`: https://arxiv.org/abs/1602.01783 +.. _`Universe Starter Agent`: https://github.com/openai/universe-starter-agent +.. _`code for this example`: https://github.com/ray-project/ray/tree/master/examples/a3c + +To run the application, first install **ray** and then some dependencies: + +.. code-block:: bash + + pip install tensorflow + pip install six + pip install gym[atari]==0.7.4 + pip install opencv-python + pip install scipy + +Note that this code **currently does not work** with ``gym==0.8.0``. + +You can run the code with + +.. code-block:: bash + + python ray/examples/a3c/driver.py [num_workers] + +Reinforcement Learning +---------------------- + +Reinforcement Learning is an area of machine learning concerned with **learning +how an agent should act in an environment** so as to maximize some form of +cumulative reward. Typically, an agent will observe the current state of the +environment and take an action based on its observation. The action will change +the state of the environment and will provide some numerical reward (or penalty) +to the agent. The agent will then take in another observation and the process +will repeat. **The mapping from state to action is a policy**, and in +reinforcement learning, this policy is often represented with a deep neural +network. + +The **environment** is often a simulator (for example, a physics engine), and +reinforcement learning algorithms often involve trying out many different +sequences of actions within these simulators. These **rollouts** can often be +done in parallel. + +Policies are often initialized randomly and incrementally improved via +simulation within the environment. To improve a policy, gradient-based updates +may be computed based on the sequences of states and actions that have been +observed. The gradient calculation is often delayed until a termination +condition is reached (that is, the simulation has finished) so that delayed +rewards have been properly accounted for. However, in the Actor Critic model, we +can begin the gradient calculation at any point in the simulation rollout by +predicting future rewards with a Value Function approximator. + +In our A3C implementation, each worker, implemented as a Ray actor, continuously +simulates the environment. The driver will create a task that runs some steps +of the simulator using the latest model, computes a gradient update, and returns +the update to the driver. Whenever a task finishes, the driver will use the +gradient update to update the model and will launch a new task with the latest +model. + +There are two main parts to the implementation - the driver and the worker. + +Worker Code Walkthrough +----------------------- + +We use a Ray Actor to simulate the environment. + +.. code-block:: python + + import numpy as np + import ray + + @ray.actor + class Runner(object): + """Actor object to start running simulation on workers. + Gradient computation is also executed on this object.""" + def __init__(self, env_name, actor_id): + # starts simulation environment, policy, and thread. + # Thread will continuously interact with the simulation environment + self.env = env = create_env(env_name) + self.id = actor_id + self.policy = LSTMPolicy() + self.runner = RunnerThread(env, self.policy, 20) + self.start() + + def start(self): + # starts the simulation thread + self.runner.start_runner() + + def pull_batch_from_queue(self): + # Implementation details removed - gets partial rollout from queue + return rollout + + def compute_gradient(self, params): + self.policy.set_weights(params) + rollout = self.pull_batch_from_queue() + batch = process_rollout(rollout, gamma=0.99, lambda_=1.0) + gradient = self.policy.get_gradients(batch) + info = {"id": self.id, + "size": len(batch.a)} + return gradient, info + +Driver Code Walkthrough +----------------------- + +The driver manages the coordination among workers and handles updating the +global model parameters. The main training script looks like the following. + + +.. code-block:: python + + import numpy as np + import ray + + def train(num_workers, env_name="PongDeterministic-v3"): + # Setup a copy of the environment + # Instantiate a copy of the policy - mainly used as a placeholder + env = create_env(env_name, None, None) + policy = LSTMPolicy(env.observation_space.shape, env.action_space.n, 0) + obs = 0 + + # Start simulations on actors + agents = [Runner(env_name, i) for i in range(num_workers)] + + # Start gradient calculation tasks on each actor + parameters = policy.get_weights() + gradient_list = [agent.compute_gradient(parameters) for agent in agents] + + while True: # Replace with your termination condition + # wait for some gradient to be computed - unblock as soon as the earliest arrives + done_id, gradient_list = ray.wait(gradient_list) + + # get the results of the task from the object store + gradient, info = ray.get(done_id)[0] + obs += info["size"] + + # apply update, get the weights from the model, start a new task on the same actor object + policy.model_update(gradient) + parameters = policy.get_weights() + gradient_list.extend([agents[info["id"]].compute_gradient(parameters)]) + return policy diff --git a/doc/source/index.rst b/doc/source/index.rst index 6b33266b7..1cbbf2141 100644 --- a/doc/source/index.rst +++ b/doc/source/index.rst @@ -29,6 +29,7 @@ learning and reinforcement learning applications.* example-hyperopt.rst example-policy-gradient.rst example-resnet.rst + example-a3c.rst example-lbfgs.md example-rl-pong.md using-ray-with-tensorflow.md diff --git a/examples/a3c/LSTM.py b/examples/a3c/LSTM.py new file mode 100644 index 000000000..bf317aa56 --- /dev/null +++ b/examples/a3c/LSTM.py @@ -0,0 +1,92 @@ +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import numpy as np +import tensorflow as tf +import tensorflow.contrib.rnn as rnn +import distutils.version +import ray +from policy import * +use_tf100_api = distutils.version.LooseVersion(tf.VERSION) >= distutils.version.LooseVersion('1.0.0') + + +class LSTMPolicy(Policy): + def setup_graph(self, ob_space, ac_space): + """Setup model used for Policy (in this A3C, both the Critic and the Actor share the model)""" + self.x = x = tf.placeholder(tf.float32, [None] + list(ob_space)) + + for i in range(4): + x = tf.nn.elu(conv2d(x, 32, "l{}".format(i + 1), [3, 3], [2, 2])) + # introduce a "fake" batch dimension of 1 after flatten so that we can do LSTM over time dim + x = tf.expand_dims(flatten(x), [0]) + + size = 256 + if use_tf100_api: + lstm = rnn.BasicLSTMCell(size, state_is_tuple=True) + else: + lstm = rnn.rnn_cell.BasicLSTMCell(size, state_is_tuple=True) + self.state_size = lstm.state_size + step_size = tf.shape(self.x)[:1] + + c_init = np.zeros((1, lstm.state_size.c), np.float32) + h_init = np.zeros((1, lstm.state_size.h), np.float32) + self.state_init = [c_init, h_init] + c_in = tf.placeholder(tf.float32, [1, lstm.state_size.c]) + h_in = tf.placeholder(tf.float32, [1, lstm.state_size.h]) + self.state_in = [c_in, h_in] + + if use_tf100_api: + state_in = rnn.LSTMStateTuple(c_in, h_in) + else: + state_in = rnn.rnn_cell.LSTMStateTuple(c_in, h_in) + lstm_outputs, lstm_state = tf.nn.dynamic_rnn( + lstm, x, initial_state=state_in, sequence_length=step_size, + time_major=False) + lstm_c, lstm_h = lstm_state + x = tf.reshape(lstm_outputs, [-1, size]) + self.logits = linear(x, ac_space, "action", normalized_columns_initializer(0.01)) + self.vf = tf.reshape(linear(x, 1, "value", normalized_columns_initializer(1.0)), [-1]) + self.state_out = [lstm_c[:1, :], lstm_h[:1, :]] + self.sample = categorical_sample(self.logits, ac_space)[0, :] + self.var_list = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, tf.get_variable_scope().name) + self.global_step = tf.get_variable("global_step", [], tf.int32, initializer=tf.constant_initializer(0, dtype=tf.int32), + trainable=False) + + def get_gradients(self, batch): + """Computing the gradient is actually model-dependent. + The LSTM needs its hidden states in order to compute the gradient accurately.""" + feed_dict = { + self.x: batch.si, + self.ac: batch.a, + self.adv: batch.adv, + self.r: batch.r, + self.state_in[0]: batch.features[0], + self.state_in[1]: batch.features[1], + } + self.local_steps += 1 + return self.sess.run(self.grads, feed_dict=feed_dict) + + def act(self, ob, c, h): + return self.sess.run([self.sample, self.vf] + self.state_out, + {self.x: [ob], self.state_in[0]: c, self.state_in[1]: h}) + + def value(self, ob, c, h): + return self.sess.run(self.vf, {self.x: [ob], self.state_in[0]: c, self.state_in[1]: h})[0] + + def get_initial_features(self): + return self.state_init + + +class RawLSTMPolicy(LSTMPolicy): + def get_weights(self): + if not hasattr(self, "_weights"): + self._weights = self.variables.get_weights() + return self._weights + + def set_weights(self, weights): + self._weights = weights + + def model_update(self, grads): + for var, grad in zip(self.var_list, grads): + self._weights[var.name[:-2]] -= 1e-4 * grad diff --git a/examples/a3c/driver.py b/examples/a3c/driver.py new file mode 100644 index 000000000..9193e1193 --- /dev/null +++ b/examples/a3c/driver.py @@ -0,0 +1,81 @@ +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import ray +import numpy as np +from runner import RunnerThread, process_rollout +from LSTM import LSTMPolicy +import tensorflow as tf +import six.moves.queue as queue +import gym +import sys +from datetime import datetime, timedelta +from misc import timestamp, Profiler +from envs import create_env + +@ray.actor +class Runner(object): + """Actor object to start running simulation on workers. + Gradient computation is also executed from this object.""" + def __init__(self, env_name, actor_id, logdir="tmp/", start=True): + env = create_env(env_name, None, None) + self.id = actor_id + num_actions = env.action_space.n + self.policy = LSTMPolicy(env.observation_space.shape, num_actions, actor_id) + self.runner = RunnerThread(env, self.policy, 20) + self.env = env + self.logdir = logdir + if start: + self.start() + + def pull_batch_from_queue(self): + """ self explanatory: take a rollout from the queue of the thread runner. """ + rollout = self.runner.queue.get(timeout=600.0) + while not rollout.terminal: + try: + rollout.extend(self.runner.queue.get_nowait()) + except queue.Empty: + break + return rollout + + def start(self): + summary_writer = tf.summary.FileWriter(self.logdir + "test_1") + self.summary_writer = summary_writer + self.runner.start_runner(self.policy.sess, summary_writer) + + def compute_gradient(self, params): + self.policy.set_weights(params) + rollout = self.pull_batch_from_queue() + batch = process_rollout(rollout, gamma=0.99, lambda_=1.0) + gradient = self.policy.get_gradients(batch) + info = {"id": self.id, + "size": len(batch.a)} + return gradient, info + + +def train(num_workers, env_name="PongDeterministic-v3"): + env = create_env(env_name, None, None) + policy = LSTMPolicy(env.observation_space.shape, env.action_space.n, 0) + agents = [Runner(env_name, i) for i in range(num_workers)] + parameters = policy.get_weights() + gradient_list = [agent.compute_gradient(parameters) for agent in agents] + steps = 0 + obs = 0 + while True: + done_id, gradient_list = ray.wait(gradient_list) + gradient, info = ray.get(done_id)[0] + policy.model_update(gradient) + parameters = policy.get_weights() + steps += 1 + obs += info["size"] + gradient_list.extend([agents[info["id"]].compute_gradient(parameters)]) + return policy + +if __name__ == '__main__': + if gym.__version__[:3] == '0.8': + raise Exception("This example currently does not work with gym==0.8.0. " + "Please downgrade to gym==0.7.4."); + NW = int(sys.argv[1]) + ray.init(num_workers=NW, num_cpus=NW) + train(NW) diff --git a/examples/a3c/envs.py b/examples/a3c/envs.py new file mode 100644 index 000000000..8a5df3492 --- /dev/null +++ b/examples/a3c/envs.py @@ -0,0 +1,166 @@ +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import cv2 +import gym +from gym.spaces.box import Box +from gym import spaces +import logging +import numpy as np +import time + +import vectorized +from vectorized.wrappers import Unvectorize, Vectorize + +logger = logging.getLogger(__name__) +logger.setLevel(logging.INFO) + +def create_env(env_id, client_id, remotes, **kwargs): + return create_atari_env(env_id) + +def create_atari_env(env_id): + env = gym.make(env_id) + env = Vectorize(env) + env = AtariRescale42x42(env) + env = DiagnosticsInfo(env) + env = Unvectorize(env) + return env + +def DiagnosticsInfo(env, *args, **kwargs): + return vectorized.VectorizeFilter(env, DiagnosticsInfoI, *args, **kwargs) + +class DiagnosticsInfoI(vectorized.Filter): + def __init__(self, log_interval=503): + super(DiagnosticsInfoI, self).__init__() + + self._episode_time = time.time() + self._last_time = time.time() + self._local_t = 0 + self._log_interval = log_interval + self._episode_reward = 0 + self._episode_length = 0 + self._all_rewards = [] + self._num_vnc_updates = 0 + self._last_episode_id = -1 + + def _after_reset(self, observation): + logger.info('Resetting environment') + self._episode_reward = 0 + self._episode_length = 0 + self._all_rewards = [] + return observation + + def _after_step(self, observation, reward, done, info): + to_log = {} + if self._episode_length == 0: + self._episode_time = time.time() + + self._local_t += 1 + if info.get("stats.vnc.updates.n") is not None: + self._num_vnc_updates += info.get("stats.vnc.updates.n") + + if self._local_t % self._log_interval == 0: + cur_time = time.time() + elapsed = cur_time - self._last_time + fps = self._log_interval / elapsed + self._last_time = cur_time + cur_episode_id = info.get('vectorized.episode_id', 0) + to_log["diagnostics/fps"] = fps + if self._last_episode_id == cur_episode_id: + to_log["diagnostics/fps_within_episode"] = fps + self._last_episode_id = cur_episode_id + if info.get("stats.gauges.diagnostics.lag.action") is not None: + to_log["diagnostics/action_lag_lb"] = info["stats.gauges.diagnostics.lag.action"][0] + to_log["diagnostics/action_lag_ub"] = info["stats.gauges.diagnostics.lag.action"][1] + if info.get("reward.count") is not None: + to_log["diagnostics/reward_count"] = info["reward.count"] + if info.get("stats.gauges.diagnostics.clock_skew") is not None: + to_log["diagnostics/clock_skew_lb"] = info["stats.gauges.diagnostics.clock_skew"][0] + to_log["diagnostics/clock_skew_ub"] = info["stats.gauges.diagnostics.clock_skew"][1] + if info.get("stats.gauges.diagnostics.lag.observation") is not None: + to_log["diagnostics/observation_lag_lb"] = info["stats.gauges.diagnostics.lag.observation"][0] + to_log["diagnostics/observation_lag_ub"] = info["stats.gauges.diagnostics.lag.observation"][1] + + if info.get("stats.vnc.updates.n") is not None: + to_log["diagnostics/vnc_updates_n"] = info["stats.vnc.updates.n"] + to_log["diagnostics/vnc_updates_n_ps"] = self._num_vnc_updates / elapsed + self._num_vnc_updates = 0 + if info.get("stats.vnc.updates.bytes") is not None: + to_log["diagnostics/vnc_updates_bytes"] = info["stats.vnc.updates.bytes"] + if info.get("stats.vnc.updates.pixels") is not None: + to_log["diagnostics/vnc_updates_pixels"] = info["stats.vnc.updates.pixels"] + if info.get("stats.vnc.updates.rectangles") is not None: + to_log["diagnostics/vnc_updates_rectangles"] = info["stats.vnc.updates.rectangles"] + if info.get("env_status.state_id") is not None: + to_log["diagnostics/env_state_id"] = info["env_status.state_id"] + + if reward is not None: + self._episode_reward += reward + if observation is not None: + self._episode_length += 1 + self._all_rewards.append(reward) + + if done: + logger.info('Episode terminating: episode_reward=%s episode_length=%s', self._episode_reward, self._episode_length) + total_time = time.time() - self._episode_time + to_log["global/episode_reward"] = self._episode_reward + to_log["global/episode_length"] = self._episode_length + to_log["global/episode_time"] = total_time + to_log["global/reward_per_time"] = self._episode_reward / total_time + self._episode_reward = 0 + self._episode_length = 0 + self._all_rewards = [] + + return observation, reward, done, to_log + +def _process_frame42(frame): + frame = frame[34:34+160, :160] + # Resize by half, then down to 42x42 (essentially mipmapping). If + # we resize directly we lose pixels that, when mapped to 42x42, + # aren't close enough to the pixel boundary. + frame = cv2.resize(frame, (80, 80)) + frame = cv2.resize(frame, (42, 42)) + frame = frame.mean(2) + frame = frame.astype(np.float32) + frame *= (1.0 / 255.0) + frame = np.reshape(frame, [42, 42, 1]) + return frame + +class AtariRescale42x42(vectorized.ObservationWrapper): + def __init__(self, env=None): + super(AtariRescale42x42, self).__init__(env) + self.observation_space = Box(0.0, 1.0, [42, 42, 1]) + + def _observation(self, observation_n): + return [_process_frame42(observation) for observation in observation_n] + + +class CropScreen(vectorized.ObservationWrapper): + """Crops out a [height]x[width] area starting from (top,left) """ + def __init__(self, env, height, width, top=0, left=0): + super(CropScreen, self).__init__(env) + self.height = height + self.width = width + self.top = top + self.left = left + self.observation_space = Box(0, 255, shape=(height, width, 3)) + + def _observation(self, observation_n): + return [ob[self.top:self.top+self.height, self.left:self.left+self.width, :] if ob is not None else None + for ob in observation_n] + +def _process_frame_flash(frame): + frame = cv2.resize(frame, (200, 128)) + frame = frame.mean(2).astype(np.float32) + frame *= (1.0 / 255.0) + frame = np.reshape(frame, [128, 200, 1]) + return frame + +class FlashRescale(vectorized.ObservationWrapper): + def __init__(self, env=None): + super(FlashRescale, self).__init__(env) + self.observation_space = Box(0.0, 1.0, [128, 200, 1]) + + def _observation(self, observation_n): + return [_process_frame_flash(observation) for observation in observation_n] diff --git a/examples/a3c/misc.py b/examples/a3c/misc.py new file mode 100644 index 000000000..99549624d --- /dev/null +++ b/examples/a3c/misc.py @@ -0,0 +1,26 @@ +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import numpy as np +from datetime import datetime +import cProfile, pstats, io + +def timestamp(): + return datetime.now().timestamp() + +class Profiler(object): + def __init__(self): + self.pr = cProfile.Profile() + pass + + def __enter__(self): + self.pr.enable() + + def __exit__(self ,type, value, traceback): + self.pr.disable() + s = io.StringIO() + sortby = 'cumtime' + ps = pstats.Stats(self.pr, stream=s).sort_stats(sortby) + ps.print_stats(.2) + print(s.getvalue()) diff --git a/examples/a3c/policy.py b/examples/a3c/policy.py new file mode 100644 index 000000000..188bbbe46 --- /dev/null +++ b/examples/a3c/policy.py @@ -0,0 +1,133 @@ +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import numpy as np +import tensorflow as tf +import tensorflow.contrib.rnn as rnn +import distutils.version +import ray +use_tf100_api = distutils.version.LooseVersion(tf.VERSION) >= distutils.version.LooseVersion('1.0.0') + +class Policy(object): + """Policy base class""" + + def __init__(self, ob_space, ac_space, task, name="local"): + self.local_steps = 0 + worker_device = "/job:localhost/replica:0/task:0/cpu:0" + self.g = tf.Graph() + with self.g.as_default(), tf.device(worker_device): + with tf.variable_scope(name): + self.setup_graph(ob_space, ac_space) + assert all([hasattr(self, attr) for attr in ["vf", "logits", "x", "var_list"]]) + print("Setting up loss") + self.setup_loss(ac_space) + self.initialize() + + def setup_graph(self): + raise NotImplementedError + + def setup_loss(self, num_actions, summarize=True): + self.ac = tf.placeholder(tf.float32, [None, num_actions], name="ac") + self.adv = tf.placeholder(tf.float32, [None], name="adv") + self.r = tf.placeholder(tf.float32, [None], name="r") + + log_prob_tf = tf.nn.log_softmax(self.logits) + prob_tf = tf.nn.softmax(self.logits) + + # the "policy gradients" loss: its derivative is precisely the policy gradient + # notice that self.ac is a placeholder that is provided externally. + # adv will contain the advantages, as calculated in process_rollout + pi_loss = - tf.reduce_sum(tf.reduce_sum(log_prob_tf * self.ac, [1]) * self.adv) + + # loss of value function + vf_loss = 0.5 * tf.reduce_sum(tf.square(self.vf - self.r)) + vf_loss = tf.Print(vf_loss, [vf_loss], "Value Fn Loss") + entropy = - tf.reduce_sum(prob_tf * log_prob_tf) + + bs = tf.to_float(tf.shape(self.x)[0]) + self.loss = pi_loss + 0.5 * vf_loss - entropy * 0.01 + + grads = tf.gradients(self.loss, self.var_list) + self.grads, _ = tf.clip_by_global_norm(grads, 40.0) + + grads_and_vars = list(zip(self.grads, self.var_list)) + opt = tf.train.AdamOptimizer(1e-4) + self._apply_gradients = opt.apply_gradients(grads_and_vars) + + if summarize: + tf.summary.scalar("model/policy_loss", pi_loss / bs) + tf.summary.scalar("model/value_loss", vf_loss / bs) + tf.summary.scalar("model/entropy", entropy / bs) + tf.summary.image("model/state", self.x) + self.summary_op = tf.summary.merge_all() + + def initialize(self): + self.sess = tf.Session(graph=self.g, config=tf.ConfigProto(intra_op_parallelism_threads=1, inter_op_parallelism_threads=2)) + self.variables = ray.experimental.TensorFlowVariables(self.loss, self.sess) + self.sess.run(tf.global_variables_initializer()) + + def model_update(self, grads): + feed_dict = {self.grads[i]: grads[i] + for i in range(len(grads))} + self.sess.run(self._apply_gradients, feed_dict=feed_dict) + + def get_weights(self): + weights = self.variables.get_weights() + return weights + + def set_weights(self, weights): + self.variables.set_weights(weights) + + def get_gradients(self, batch): + raise NotImplementedError + + def get_vf_loss(self): + raise NotImplementedError + + + def act(self, ob): + raise NotImplementedError + + def value(self, ob): + raise NotImplementedError + +def normalized_columns_initializer(std=1.0): + def _initializer(shape, dtype=None, partition_info=None): + out = np.random.randn(*shape).astype(np.float32) + out *= std / np.sqrt(np.square(out).sum(axis=0, keepdims=True)) + return tf.constant(out) + return _initializer + +def flatten(x): + return tf.reshape(x, [-1, np.prod(x.get_shape().as_list()[1:])]) + +def conv2d(x, num_filters, name, filter_size=(3, 3), stride=(1, 1), pad="SAME", dtype=tf.float32, collections=None): + with tf.variable_scope(name): + stride_shape = [1, stride[0], stride[1], 1] + filter_shape = [filter_size[0], filter_size[1], int(x.get_shape()[3]), num_filters] + + # there are "num input feature maps * filter height * filter width" + # inputs to each hidden unit + fan_in = np.prod(filter_shape[:3]) + # each unit in the lower layer receives a gradient from: + # "num output feature maps * filter height * filter width" / + # pooling size + fan_out = np.prod(filter_shape[:2]) * num_filters + # initialize weights with random weights + w_bound = np.sqrt(6. / (fan_in + fan_out)) + + w = tf.get_variable("W", filter_shape, dtype, tf.random_uniform_initializer(-w_bound, w_bound), + collections=collections) + b = tf.get_variable("b", [1, 1, 1, num_filters], initializer=tf.constant_initializer(0.0), + collections=collections) + return tf.nn.conv2d(x, w, stride_shape, pad) + b + +def linear(x, size, name, initializer=None, bias_init=0): + w = tf.get_variable(name + "/w", [x.get_shape()[1], size], initializer=initializer) + b = tf.get_variable(name + "/b", [size], initializer=tf.constant_initializer(bias_init)) + return tf.matmul(x, w) + b + +def categorical_sample(logits, d): + value = tf.squeeze(tf.multinomial(logits - tf.reduce_max(logits, [1], keep_dims=True), 1), [1]) + return tf.one_hot(value, d) diff --git a/examples/a3c/runner.py b/examples/a3c/runner.py new file mode 100644 index 000000000..e43b10b5d --- /dev/null +++ b/examples/a3c/runner.py @@ -0,0 +1,159 @@ +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from collections import namedtuple +import numpy as np +import tensorflow as tf +from LSTM import LSTMPolicy +import six.moves.queue as queue +import scipy.signal +import threading +import distutils.version +use_tf12_api = distutils.version.LooseVersion(tf.VERSION) >= distutils.version.LooseVersion('0.12.0') +from datetime import datetime + +def discount(x, gamma): + return scipy.signal.lfilter([1], [1, -gamma], x[::-1], axis=0)[::-1] + +def process_rollout(rollout, gamma, lambda_=1.0): + """ +given a rollout, compute its returns and the advantage +""" + batch_si = np.asarray(rollout.states) + batch_a = np.asarray(rollout.actions) + rewards = np.asarray(rollout.rewards) + vpred_t = np.asarray(rollout.values + [rollout.r]) + + rewards_plus_v = np.asarray(rollout.rewards + [rollout.r]) + batch_r = discount(rewards_plus_v, gamma)[:-1] + delta_t = rewards + gamma * vpred_t[1:] - vpred_t[:-1] + # this formula for the advantage comes "Generalized Advantage Estimation": + # https://arxiv.org/abs/1506.02438 + batch_adv = discount(delta_t, gamma * lambda_) + + features = rollout.features[0] + return Batch(batch_si, batch_a, batch_adv, batch_r, rollout.terminal, features) + +Batch = namedtuple("Batch", ["si", "a", "adv", "r", "terminal", "features"]) + +class PartialRollout(object): + """ +a piece of a complete rollout. We run our agent, and process its experience +once it has processed enough steps. +""" + def __init__(self): + self.states = [] + self.actions = [] + self.rewards = [] + self.values = [] + self.r = 0.0 + self.terminal = False + self.features = [] + + def add(self, state, action, reward, value, terminal, features): + self.states += [state] + self.actions += [action] + self.rewards += [reward] + self.values += [value] + self.terminal = terminal + self.features += [features] + + def extend(self, other): + assert not self.terminal + self.states.extend(other.states) + self.actions.extend(other.actions) + self.rewards.extend(other.rewards) + self.values.extend(other.values) + self.r = other.r + self.terminal = other.terminal + self.features.extend(other.features) + +class RunnerThread(threading.Thread): + """ This thread constantly interacts with the environment and tell it what to do. """ + def __init__(self, env, policy, num_local_steps, visualise=False): + threading.Thread.__init__(self) + self.queue = queue.Queue(5) + self.num_local_steps = num_local_steps + self.env = env + self.last_features = None + self.policy = policy + self.daemon = True + self.sess = None + self.summary_writer = None + self.visualise = visualise + + def start_runner(self, sess, summary_writer): + self.sess = sess + self.summary_writer = summary_writer + self.start() + + def run(self): + with self.sess.as_default(): + self._run() + + def _run(self): + rollout_provider = env_runner(self.env, self.policy, self.num_local_steps, self.summary_writer, self.visualise) + while True: + # the timeout variable exists because apparently, if one worker dies, the other workers + # won't die with it, unless the timeout is set to some large number. This is an empirical + # observation. + self.queue.put(next(rollout_provider), timeout=600.0) + # print("Current Q count: %d" % len(self.queue.queue)) + + + +def env_runner(env, policy, num_local_steps, summary_writer, render): + """ +The logic of the thread runner. In brief, it constantly keeps on running +the policy, and as long as the rollout exceeds a certain length, the thread +runner appends the policy to the queue. +""" + last_state = env.reset() + last_features = policy.get_initial_features() + length = 0 + rewards = 0 + + while True: + terminal_end = False + rollout = PartialRollout() + + for _ in range(num_local_steps): + fetched = policy.act(last_state, *last_features) + action, value_, features = fetched[0], fetched[1], fetched[2:] + # argmax to convert from one-hot + state, reward, terminal, info = env.step(action.argmax()) + if render: + env.render() + + # collect the experience + rollout.add(last_state, action, reward, value_, terminal, last_features) + length += 1 + rewards += reward + + last_state = state + last_features = features + + if info: + summary = tf.Summary() + for k, v in info.items(): + summary.value.add(tag=k, simple_value=float(v)) + summary_writer.add_summary(summary, policy.global_step.eval()) + summary_writer.flush() + + timestep_limit = env.spec.tags.get('wrapper_config.TimeLimit.max_episode_steps') + if terminal or length >= timestep_limit: + terminal_end = True + if length >= timestep_limit or not env.metadata.get('semantics.autoreset'): + last_state = env.reset() + last_features = policy.get_initial_features() + # print("Episode finished. Sum of rewards: %d. Length: %d" % (rewards, length)) + length = 0 + rewards = 0 + break + + if not terminal_end: + rollout.r = policy.value(last_state, *last_features) + + # once we have enough experience, yield it, and have the ThreadRunner place it on a queue + yield rollout diff --git a/examples/a3c/vectorized/__init__.py b/examples/a3c/vectorized/__init__.py new file mode 100644 index 000000000..6e135ba67 --- /dev/null +++ b/examples/a3c/vectorized/__init__.py @@ -0,0 +1,8 @@ +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from .vectorize_core import Env, Wrapper, ObservationWrapper, ActionWrapper, RewardWrapper +from .multiprocessing_env import MultiprocessingEnv +from .vectorize_filter import Filter, VectorizeFilter +from .wrappers import Vectorize, Unvectorize, WeakUnvectorize diff --git a/examples/a3c/vectorized/monitoring.py b/examples/a3c/vectorized/monitoring.py new file mode 100644 index 000000000..1b3f54fe5 --- /dev/null +++ b/examples/a3c/vectorized/monitoring.py @@ -0,0 +1,56 @@ +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import weakref + +from gym import monitoring + +class Monitor(object): + def __init__(self, env_n): + """env_n is a collection of unvectorized envs""" + self.monitor_n = [monitoring.Monitor(env) for env in env_n] + + @property + def env(self): + # The real env is the first unwrapped env. Maybe we should + # maintain our own weakref rather than doing this. + return self.monitor_n[0].env.env + + def start(self, directory, video_callable=None, seed_n=None, force=False, + resume=False, write_upon_reset=False, uid=None): + if seed_n is None: + seed_n = [None] * len(self.monitor_n) + # There's way to seed just one of the vectorized environments, + # so we have to do the seeding ourselves outside of the + # underlying monitor instances. + # + # The monitor will call the .seed method on the + # WeakUnvectorized env, which just returns rather than + # actually re-seeding the env. + self.env.seed(seed_n) + + for i, monitor in enumerate(self.monitor_n): + # Only allow recording of video in first monitor + if i > 0: + video_callable = False + # Seed gets passed in but just recorded, not used. + monitor.start(directory=directory, video_callable=video_callable, + force=force, resume=resume, write_upon_reset=write_upon_reset, uid=uid) + + def close(self, *args, **kwargs): + [monitor.close(*args, **kwargs) for monitor in self.monitor_n] + + def _before_reset(self): + return [monitor._before_reset() for monitor in self.monitor_n] + + def _after_reset(self, observation_n): + assert len(observation_n) == len(self.monitor_n) + return [monitor._after_reset(observation) for monitor, observation in zip(self.monitor_n, observation_n)] + + def _before_step(self, action_n): + assert len(action_n) == len(self.monitor_n) + return [monitor._before_step(action) for monitor, action in zip(self.monitor_n, action_n)] + + def _after_step(self, observation_n, reward_n, done_n, info): + return [monitor._after_step(o, r, d, i) for monitor, o, r, d, i in zip(self.monitor_n, observation_n, reward_n, done_n, info['n'])] diff --git a/examples/a3c/vectorized/multiprocessing_env.py b/examples/a3c/vectorized/multiprocessing_env.py new file mode 100644 index 000000000..98cc78942 --- /dev/null +++ b/examples/a3c/vectorized/multiprocessing_env.py @@ -0,0 +1,326 @@ +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import logging +import multiprocessing +import numpy as np +import traceback + +import gym +from gym import spaces +import vectorized.vectorize_core as core + +logger = logging.getLogger(__name__) +logger.setLevel(logging.INFO) + +class Error(Exception): + pass + +def display_name(exception): + prefix = '' + # AttributeError has no __module__; RuntimeError has module of + # exceptions + if hasattr(exception, '__module__') and exception.__module__ != 'exceptions': + prefix = exception.__module__ + '.' + return prefix + type(exception).__name__ + +def render_dict(error): + return { + 'type': display_name(error), + 'message': error.message, + 'traceback': traceback.format_exc(error) + } + +class Worker(object): + def __init__(self, env_m, worker_idx): + # These are instantiated in the *parent* process + # currently. Probably will want to change this. The parent + # does need to obtain the relevant Spaces at some stage, but + # that's doable. + self.worker_idx = worker_idx + self.env_m = env_m + self.m = len(env_m) + self.parent_conn, self.child_conn = multiprocessing.Pipe() + self.joiner = multiprocessing.Process(target=self.run) + self._clear_state() + + self.start() + + # Parent only! + self.child_conn.close() + + def _clear_state(self): + self.mask = [True] * self.m + + # Control methods + + def start(self): + self.joiner.start() + + def _parent_recv(self): + rendered, res = self.parent_conn.recv() + if rendered is not None: + raise Error('[Worker {}] Error: {} ({})\n\n{}'.format(self.worker_idx, rendered['message'], rendered['type'], rendered['traceback'])) + return res + + def _child_send(self, msg): + self.child_conn.send((None, msg)) + + def _parent_send(self, msg): + try: + self.parent_conn.send(msg) + except IOError: # the worker is now dead + try: + res = self._parent_recv() + except EOFError: + raise Error('[Worker {}] Child died unexpectedly'.format(self.worker_idx)) + else: + raise Error('[Worker {}] Child returned unexpected result: {}'.format(self.worker_idx, res)) + + def close_start(self): + self._parent_send(('close', None)) + + def close_finish(self): + self.joiner.join() + + def reset_start(self): + self._parent_send(('reset', None)) + + def reset_finish(self): + return self._parent_recv() + + def step_start(self, action_m): + """action_m: the batch of actions for this worker""" + self._parent_send(('step', action_m)) + + def step_finish(self): + return self._parent_recv() + + def mask_start(self, i): + self._parent_send(('mask', i)) + + def seed_start(self, seed_m): + self._parent_send(('seed', seed_m)) + + def render_start(self, mode, close): + self._parent_send(('render', (mode, close))) + + def render_finish(self): + return self._parent_recv() + + def run(self): + try: + self.do_run() + except Exception as e: + rendered = render_dict(e) + self.child_conn.send((rendered, None)) + return + + def do_run(self): + # Child only! + self.parent_conn.close() + + while True: + method, body = self.child_conn.recv() + logger.debug('[%d] Received: method=%s body=%s', self.worker_idx, method, body) + if method == 'close': + logger.info('Closing envs') + # TODO: close envs? + return + elif method == 'reset': + self._clear_state() + observation_m = [env.reset() for env in self.env_m] + self._child_send(observation_m) + elif method == 'step': + action_m = body + observation_m, reward_m, done_m, info = self.step_m(action_m) + self._child_send((observation_m, reward_m, done_m, info)) + elif method == 'mask': + i = body + assert 0 <= i < self.m, 'Bad value for mask: {} (should be >= 0 and < {})'.format(i, self.m) + + self.mask[i] = False + logger.debug('[%d] Applying mask: i=%d', self.worker_idx, i) + elif method == 'seed': + seeds = body + [env.seed(seed) for env, seed in zip(self.env_m, seeds)] + elif method == 'render': + mode, close = body + if mode == 'human': + self.env_m[0].render(mode=mode, close=close) + result = [None] + else: + result = [env.render(mode=mode, close=close) for env in self.env_m] + self._child_send(result) + else: + raise Error('Bad method: {}'.format(method)) + + def step_m(self, action_m): + observation_m = [] + reward_m = [] + done_m = [] + info = {'m': []} + + for env, enabled, action in zip(self.env_m, self.mask, action_m): + if enabled: + observation, reward, done, info_i = env.step(action) + if done: + observation = env.reset() + else: + observation = None + reward = 0 + done = False + info_i = {} + observation_m.append(observation) + reward_m.append(reward) + done_m.append(done) + info['m'].append(info_i) + return observation_m, reward_m, done_m, info + + +def step_n(worker_n, action_n): + accumulated = 0 + for worker in worker_n: + action_m = action_n[accumulated:accumulated+worker.m] + worker.step_start(action_m) + accumulated += worker.m + + observation_n = [] + reward_n = [] + done_n = [] + info = {'n': []} + + for worker in worker_n: + observation_m, reward_m, done_m, info_i = worker.step_finish() + observation_n += observation_m + reward_n += reward_m + done_n += done_m + info['n'] += info_i['m'] + return observation_n, reward_n, done_n, info + + +def reset_n(worker_n): + for worker in worker_n: + worker.reset_start() + + observation_n = [] + for worker in worker_n: + observation_n += worker.reset_finish() + + return observation_n + + +def seed_n(worker_n, seed_n): + accumulated = 0 + for worker in worker_n: + action_m = seed_n[accumulated:accumulated+worker.m] + worker.seed_start(seed_n) + accumulated += worker.m + + +def mask(worker_n, i): + accumulated = 0 + for k, worker in enumerate(worker_n): + if accumulated + worker.m <= i: + accumulated += worker.m + else: + worker.mask_start(i - accumulated) + return + +def render_n(worker_n, mode, close): + if mode == 'human': + # Only render 1 worker + worker_n = worker_n[0:] + + for worker in worker_n: + worker.render_start(mode, close) + res = [] + for worker in worker_n: + res += worker.render_finish() + if mode != 'human': + return res + else: + return None + +def close_n(worker_n): + if worker_n is None: + return + + # TODO: better error handling: workers should die when we go away + # anyway. Also technically should wait for these processes if + # we're not crashing. + for worker in worker_n: + try: + worker.close_start() + except Error: + pass + + # for worker in worker_n: + # try: + # worker.close_finish() + # except Error: + # pass + +class MultiprocessingEnv(core.Env): + metadata = { + 'runtime.vectorized': True, + } + + def __init__(self, env_id): + self.worker_n = None + + # Pull the relevant info from a transient env instance + self.spec = gym.spec(env_id) + env = self.spec.make() + + current_metadata = self.metadata + self.metadata = env.metadata.copy() + self.metadata.update(current_metadata) + + self.action_space = env.action_space + self.observation_space = env.observation_space + self.reward_range = env.reward_range + + def _configure(self, n=1, pool_size=None, episode_limit=None): + super(MultiprocessingEnv, self)._configure() + self.n = n + self.envs = [self.spec.make() for _ in range(self.n)] + + if pool_size is None: + pool_size = min(len(self.envs), multiprocessing.cpu_count() - 1) + pool_size = max(1, pool_size) + + self.worker_n = [] + m = int((self.n + pool_size - 1) / pool_size) + for i in range(0, self.n, m): + envs = self.envs[i:i+m] + self.worker_n.append(Worker(envs, i)) + + if episode_limit is not None: + self._episode_id.episode_limit = episode_limit + + def _seed(self, seed): + seed_n(self.worker_n, seed) + return [[seed_i] for seed_i in seed] + + def _reset(self): + return reset_n(self.worker_n) + + def _step(self, action_n): + return step_n(self.worker_n, action_n) + + def _render(self, mode='human', close=False): + return render_n(self.worker_n, mode=mode, close=close) + + def mask(self, i): + mask(self.worker_n, i) + + def _close(self): + close_n(self.worker_n) + +if __name__ == '__main__': + env_n = make('Pong-v3') + env_n.configure() + env_n.reset() + print(env_n.step([0] * 10)) diff --git a/examples/a3c/vectorized/vectorize_core.py b/examples/a3c/vectorized/vectorize_core.py new file mode 100644 index 000000000..aa1ed4c46 --- /dev/null +++ b/examples/a3c/vectorized/vectorize_core.py @@ -0,0 +1,53 @@ +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import gym +from gym import spaces + +class Env(gym.Env): + """Base class capable of handling vectorized environments. + """ + metadata = { + # This key indicates whether an env is vectorized (or, in the case of + # Wrappers where autovectorize=True, whether they should automatically + # be wrapped by a Vectorize wrapper.) + 'runtime.vectorized': True, + } + + # Number of remotes. User should set this. + n = None + +class Wrapper(Env, gym.Wrapper): + """Use this instead of gym.Wrapper iff you're wrapping a vectorized env, + (or a vanilla env you wish to be vectorized). + """ + # If True and this is instantiated with a non-vectorized environment, + # automatically wrap it with the Vectorize wrapper. + autovectorize = True + + def __init__(self, env): + super(Wrapper, self).__init__(env) + if not env.metadata.get('runtime.vectorized'): + if self.autovectorize: + # Circular dependency :( + import vectorize.wrappers as wrappers + env = wrappers.Vectorize(env) + else: + raise Exception('This wrapper can only wrap vectorized envs (i.e. where env.metadata["runtime.vectorized"] = True), not {}. Set "self.autovectorize = True" to automatically add a Vectorize wrapper.'.format(env)) + + self.env = env + + def _configure(self, **kwargs): + super(Wrapper, self)._configure(**kwargs) + assert self.env.n is not None, "Did not set self.env.n: self.n={} self.env={} self={}".format(self.env.n, self.env, self) + self.n = self.env.n + +class ObservationWrapper(Wrapper, gym.ObservationWrapper): + pass + +class RewardWrapper(Wrapper, gym.RewardWrapper): + pass + +class ActionWrapper(Wrapper, gym.ActionWrapper): + pass \ No newline at end of file diff --git a/examples/a3c/vectorized/vectorize_filter.py b/examples/a3c/vectorized/vectorize_filter.py new file mode 100644 index 000000000..eccc477e3 --- /dev/null +++ b/examples/a3c/vectorized/vectorize_filter.py @@ -0,0 +1,54 @@ +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import vectorized.vectorize_core as core + +class Filter(object): + def _after_reset(self, observation): + return observation + + def _after_step(self, observation, reward, done, info): + return observation, reward, done, info + +class VectorizeFilter(core.Wrapper): + """Vectorizes a Filter written for the non-vectorized case.""" + + autovectorize = False + metadata = { + 'configure.required': True + } + + def __init__(self, env, filter_factory, *args, **kwargs): + super(VectorizeFilter, self).__init__(env) + self.filter_factory = filter_factory + self._args = args + self._kwargs = kwargs + + def _configure(self, **kwargs): + super(VectorizeFilter, self)._configure(**kwargs) + self.filter_n = [self.filter_factory(*self._args, **self._kwargs) for _ in range(self.n)] + + def _reset(self): + observation_n = self.env.reset() + observation_n = [filter._after_reset(observation) for filter, observation in zip(self.filter_n, observation_n)] + return observation_n + + def _step(self, action_n): + o_n, r_n, d_n, i = self.env.step(action_n) + + observation_n = [] + reward_n = [] + done_n = [] + info = i.copy() + info['n'] = [] + for filter, observation, reward, done, info_i in zip(self.filter_n, o_n, r_n, d_n, i['n']): + observation, reward, done, info_i = filter._after_step(observation, reward, done, info_i) + observation_n.append(observation) + reward_n.append(reward) + done_n.append(done) + info['n'].append(info_i) + return observation_n, reward_n, done_n, info + + def __str__(self): + return '<{}[{}]{}>'.format(type(self).__name__, self.filter_factory, self.env) diff --git a/examples/a3c/vectorized/wrappers.py b/examples/a3c/vectorized/wrappers.py new file mode 100644 index 000000000..83058cf04 --- /dev/null +++ b/examples/a3c/vectorized/wrappers.py @@ -0,0 +1,93 @@ +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import gym +import weakref + +import vectorized.vectorize_core as core + +class Vectorize(gym.Wrapper): + """ +Given an unvectorized environment (where, e.g., the output of .step() is an observation +rather than a list of observations), turn it into a vectorized environment with a batch of size +1. +""" + + metadata = {'runtime.vectorized': True} + + def __init__(self, env): + super(Vectorize, self).__init__(env) + assert not env.metadata.get('runtime.vectorized') + assert self.metadata.get('runtime.vectorized') + self.n = 1 + + def _reset(self): + observation = self.env.reset() + return [observation] + + def _step(self, action): + observation, reward, done, info = self.env.step(action[0]) + return [observation], [reward], [done], {'n': [info]} + + def _seed(self, seed): + return [self.env.seed(seed[0])] + +class Unvectorize(core.Wrapper): + """ +Take a vectorized environment with a batch of size 1 and turn it into an unvectorized environment. +""" + autovectorize = False + metadata = {'runtime.vectorized': False} + + def _configure(self, **kwargs): + super(Unvectorize, self)._configure(**kwargs) + if self.n != 1: + raise Exception('Can only disable vectorization with n=1, not n={}'.format(self.n)) + + def _reset(self): + observation_n = self.env.reset() + return observation_n[0] + + def _step(self, action): + action_n = [action] + observation_n, reward_n, done_n, info = self.env.step(action_n) + return observation_n[0], reward_n[0], done_n[0], info['n'][0] + + def _seed(self, seed): + return self.env.seed([seed])[0] + + +class WeakUnvectorize(Unvectorize): + def __init__(self, env, i): + self._env_ref = weakref.ref(env) + super(WeakUnvectorize, self).__init__(env) + # WeakUnvectorize won't get configure called on it + self.i = i + + def _check_for_duplicate_wrappers(self): + pass # Disable this check because we need to wrap vectorized envs in multiple unvectorize wrappers + + @property + def env(self): + # Called upon instantiation + if not hasattr(self, '_env_ref'): + return + + env = self._env_ref() + if env is None: + raise Exception("env has been garbage collected. To keep using WeakUnvectorize, you must keep around a reference to the env object. (HINT: try assigning the env to a variable in your code.)") + return env + + @env.setter + def env(self, value): + # We'll maintain our own weakref, thank you very much. + pass + + def _seed(self, seed): + # We handle the seeding ourselves in the vectorized Monitor + return [seed] + + def close(self): + # Don't want to close through this wrapper + pass