Initial A3C Example - PongDeterministic-v3 (#331)

* Initializing A3C code

* Modifications for Ray usage

* cleanup

* removing universe dependency

* fixes (not yet working

* hack

* documentation

* Cleanup

* Preliminary Portion

Make sure to change when merging

* RL part

* Cleaning up Driver and Worker code

* Updating driver code

* instructions...

* fixed

* Minor changes.

* Fixing cmake issues

* ray instruction

* updating port to new universe

* Fix for env.configure

* redundant commands

* Revert scipy.misc -> cv2 and raise exception for wrong gym version.

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

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

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

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)

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]

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

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)

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

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

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'])]

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

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

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)

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