Package the project

deep_rl/component/__init__.py

@@ -0,0 +1,6 @@
+from .atari_wrapper import  *
+from .policy import *
+from .replay import *
+from .task import *
+from .random_process import *
+from .bench import *

deep_rl/component/atari_wrapper.py

@@ -0,0 +1,264 @@
+# based on https://github.com/openai/baselines/blob/master/baselines/common/atari_wrappers.py
+
+import numpy as np
+import gym
+from gym import spaces
+from gym.spaces import Box
+import cv2
+cv2.ocl.setUseOpenCL(False)
+
+class NoopResetEnv(gym.Wrapper):
+    def __init__(self, env, noop_max=30):
+        """Sample initial states by taking random number of no-ops on reset.
+        No-op is assumed to be action 0.
+        """
+        gym.Wrapper.__init__(self, env)
+        self.noop_max = noop_max
+        self.override_num_noops = None
+        self.noop_action = 0
+        assert env.unwrapped.get_action_meanings()[0] == 'NOOP'
+
+    def reset(self, **kwargs):
+        """ Do no-op action for a number of steps in [1, noop_max]."""
+        self.env.reset(**kwargs)
+        if self.override_num_noops is not None:
+            noops = self.override_num_noops
+        else:
+            noops = self.unwrapped.np_random.randint(1, self.noop_max + 1) #pylint: disable=E1101
+        assert noops > 0
+        obs = None
+        for _ in range(noops):
+            obs, _, done, _ = self.env.step(self.noop_action)
+            if done:
+                obs = self.env.reset(**kwargs)
+        return obs
+
+    def step(self, ac):
+        return self.env.step(ac)
+
+class FireResetEnv(gym.Wrapper):
+    def __init__(self, env):
+        """Take action on reset for environments that are fixed until firing."""
+        gym.Wrapper.__init__(self, env)
+        assert env.unwrapped.get_action_meanings()[1] == 'FIRE'
+        assert len(env.unwrapped.get_action_meanings()) >= 3
+
+    def reset(self, **kwargs):
+        self.env.reset(**kwargs)
+        obs, _, done, _ = self.env.step(1)
+        if done:
+            self.env.reset(**kwargs)
+        obs, _, done, _ = self.env.step(2)
+        if done:
+            self.env.reset(**kwargs)
+        return obs
+
+    def step(self, ac):
+        return self.env.step(ac)
+
+class EpisodicLifeEnv(gym.Wrapper):
+    def __init__(self, env):
+        """Make end-of-life == end-of-episode, but only reset on true game over.
+        Done by DeepMind for the DQN and co. since it helps value estimation.
+        """
+        gym.Wrapper.__init__(self, env)
+        self.lives = 0
+        self.was_real_done  = True
+
+    def step(self, action):
+        obs, reward, done, info = self.env.step(action)
+        self.was_real_done = done
+        # check current lives, make loss of life terminal,
+        # then update lives to handle bonus lives
+        lives = self.env.unwrapped.ale.lives()
+        if lives < self.lives and lives > 0:
+            # for Qbert sometimes we stay in lives == 0 condtion for a few frames
+            # so its important to keep lives > 0, so that we only reset once
+            # the environment advertises done.
+            done = True
+        self.lives = lives
+        return obs, reward, done, info
+
+    def reset(self, **kwargs):
+        """Reset only when lives are exhausted.
+        This way all states are still reachable even though lives are episodic,
+        and the learner need not know about any of this behind-the-scenes.
+        """
+        if self.was_real_done:
+            obs = self.env.reset(**kwargs)
+        else:
+            # no-op step to advance from terminal/lost life state
+            obs, _, _, _ = self.env.step(0)
+        self.lives = self.env.unwrapped.ale.lives()
+        return obs
+
+class MaxAndSkipEnv(gym.Wrapper):
+    def __init__(self, env, skip=4):
+        """Return only every `skip`-th frame"""
+        gym.Wrapper.__init__(self, env)
+        # most recent raw observations (for max pooling across time steps)
+        self._obs_buffer = np.zeros((2,)+env.observation_space.shape, dtype=np.uint8)
+        self._skip       = skip
+
+    def step(self, action):
+        """Repeat action, sum reward, and max over last observations."""
+        total_reward = 0.0
+        done = None
+        for i in range(self._skip):
+            obs, reward, done, info = self.env.step(action)
+            if i == self._skip - 2: self._obs_buffer[0] = obs
+            if i == self._skip - 1: self._obs_buffer[1] = obs
+            total_reward += reward
+            if done:
+                break
+        # Note that the observation on the done=True frame
+        # doesn't matter
+        max_frame = self._obs_buffer.max(axis=0)
+
+        return max_frame, total_reward, done, info
+
+    def reset(self, **kwargs):
+        return self.env.reset(**kwargs)
+
+class SkipEnv(gym.Wrapper):
+    def __init__(self, env=None, skip=4):
+        """Return only every `skip`-th frame"""
+        super(SkipEnv, self).__init__(env)
+        self._skip       = skip
+
+    def step(self, action):
+        total_reward = 0.0
+        done = None
+        for _ in range(self._skip):
+            obs, reward, done, info = self.env.step(action)
+            total_reward += reward
+            if done:
+                break
+
+        return obs, total_reward, done, info
+
+    def reset(self):
+        obs = self.env.reset()
+        return obs
+
+class WarpFrame(gym.ObservationWrapper):
+    def __init__(self, env):
+        """Warp frames to 84x84 as done in the Nature paper and later work."""
+        gym.ObservationWrapper.__init__(self, env)
+        self.width = 84
+        self.height = 84
+        self.observation_space = spaces.Box(low=0, high=255,
+            shape=(self.height, self.width, 1), dtype=np.uint8)
+
+    def observation(self, frame):
+        frame = cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY)
+        frame = cv2.resize(frame, (self.width, self.height), interpolation=cv2.INTER_AREA)
+        return frame[:, :, None]
+
+class LazyFrames(object):
+    def __init__(self, frames):
+        """This object ensures that common frames between the observations are only stored once.
+        It exists purely to optimize memory usage which can be huge for DQN's 1M frames replay
+        buffers.
+
+        This object should only be converted to numpy array before being passed to the model.
+
+        You'd not believe how complex the previous solution was."""
+        self._frames = frames
+        self._out = None
+
+    def _force(self):
+        if self._out is None:
+            self._out = np.concatenate(self._frames, axis=2)
+            self._frames = None
+        return self._out
+
+    def __array__(self, dtype=None):
+        out = self._force()
+        if dtype is not None:
+            out = out.astype(dtype)
+        return out
+
+    def __len__(self):
+        return len(self._force())
+
+    def __getitem__(self, i):
+        return self._force()[i]
+
+class StackFrame(gym.Wrapper):
+    def __init__(self, env=None, history_length=1):
+        super(StackFrame, self).__init__(env)
+        self.history_length = history_length
+        self.buffer = None
+
+    def reset(self):
+        state = self.env.reset()
+        self.buffer = [state] * self.history_length
+        # return LazyFrames(self.buffer)
+        return np.asarray(np.vstack(self.buffer))
+
+    def step(self, action):
+        state, reward, done, info = self.env.step(action)
+        self.buffer.pop(0)
+        self.buffer.append(state)
+        # return LazyFrames(self.buffer), reward, done, info
+        return np.asarray(np.vstack(self.buffer)), reward, done, info
+
+class WrapPyTorch(gym.ObservationWrapper):
+    # from https://github.com/ikostrikov/pytorch-a2c-ppo-acktr/blob/master/envs.py
+    def __init__(self, env=None):
+        super(WrapPyTorch, self).__init__(env)
+        obs_shape = self.observation_space.shape
+        self.observation_space = Box(
+            self.observation_space.low[0,0,0],
+            self.observation_space.high[0,0,0],
+            [obs_shape[2], obs_shape[1], obs_shape[0]],
+            dtype=np.uint8
+        )
+
+    def observation(self, observation):
+        return observation.transpose(2, 0, 1)
+
+class DatasetEnv(gym.Wrapper):
+    def __init__(self, env=None):
+        super(DatasetEnv, self).__init__(env)
+        self.saved_obs = []
+        self.saved_actions = []
+
+    def get_saved(self):
+        return self.saved_obs, self.saved_actions
+
+    def clear_saved(self):
+        self.saved_obs = []
+        self.saved_actions = []
+
+    def step(self, action):
+        obs, reward, done, info = self.env.step(action)
+        self.saved_actions.append(action)
+        self.saved_obs.append(obs)
+        return obs, reward, done, info
+
+    def reset(self):
+        obs = self.env.reset()
+        self.saved_obs.append(obs)
+        return obs
+
+def make_atari(env_id, frame_skip=4):
+    env = gym.make(env_id)
+    assert 'NoFrameskip' in env.spec.id
+    env = NoopResetEnv(env, noop_max=30)
+    env = MaxAndSkipEnv(env, skip=4)
+    return env
+
+def wrap_deepmind(env, episode_life=True, history_length=0):
+    """Configure environment for DeepMind-style Atari.
+    """
+    if episode_life:
+        env = EpisodicLifeEnv(env)
+    if 'FIRE' in env.unwrapped.get_action_meanings():
+        env = FireResetEnv(env)
+    env = WarpFrame(env)
+    env = WrapPyTorch(env)
+    if history_length:
+        env = StackFrame(env, history_length)
+    return env

deep_rl/component/bench.py

@@ -0,0 +1,136 @@
+# from https://raw.githubusercontent.com/openai/baselines/master/baselines/bench/monitor.py
+
+import gym
+from gym.core import Wrapper
+import time
+from glob import glob
+import csv
+import os.path as osp
+import json
+
+class Monitor(Wrapper):
+    EXT = "monitor.csv"
+    f = None
+
+    def __init__(self, env, filename, allow_early_resets=True, reset_keywords=(), info_keywords=()):
+        Wrapper.__init__(self, env=env)
+        self.tstart = time.time()
+        if filename is None:
+            self.f = None
+            self.logger = None
+        else:
+            if not filename.endswith(Monitor.EXT):
+                if osp.isdir(filename):
+                    filename = osp.join(filename, Monitor.EXT)
+                else:
+                    filename = filename + "." + Monitor.EXT
+            self.f = open(filename, "wt")
+            self.f.write('#%s\n'%json.dumps({"t_start": self.tstart, 'env_id' : env.spec and env.spec.id}))
+            self.logger = csv.DictWriter(self.f, fieldnames=('r', 'l', 't')+reset_keywords+info_keywords)
+            self.logger.writeheader()
+            self.f.flush()
+
+        self.reset_keywords = reset_keywords
+        self.info_keywords = info_keywords
+        self.allow_early_resets = allow_early_resets
+        self.rewards = None
+        self.needs_reset = True
+        self.episode_rewards = []
+        self.episode_lengths = []
+        self.episode_times = []
+        self.total_steps = 0
+        self.current_reset_info = {} # extra info about the current episode, that was passed in during reset()
+
+    def reset(self, **kwargs):
+        if not self.allow_early_resets and not self.needs_reset:
+            raise RuntimeError("Tried to reset an environment before done. If you want to allow early resets, wrap your env with Monitor(env, path, allow_early_resets=True)")
+        self.rewards = []
+        self.needs_reset = False
+        for k in self.reset_keywords:
+            v = kwargs.get(k)
+            if v is None:
+                raise ValueError('Expected you to pass kwarg %s into reset'%k)
+            self.current_reset_info[k] = v
+        return self.env.reset(**kwargs)
+
+    def step(self, action):
+        if self.needs_reset:
+            raise RuntimeError("Tried to step environment that needs reset")
+        ob, rew, done, info = self.env.step(action)
+        self.rewards.append(rew)
+        if done:
+            self.needs_reset = True
+            eprew = sum(self.rewards)
+            eplen = len(self.rewards)
+            epinfo = {"r": round(eprew, 6), "l": eplen, "t": round(time.time() - self.tstart, 6)}
+            for k in self.info_keywords:
+                epinfo[k] = info[k]
+            self.episode_rewards.append(eprew)
+            self.episode_lengths.append(eplen)
+            self.episode_times.append(time.time() - self.tstart)
+            epinfo.update(self.current_reset_info)
+            if self.logger:
+                self.logger.writerow(epinfo)
+                self.f.flush()
+            info['episode'] = epinfo
+        self.total_steps += 1
+        return (ob, rew, done, info)
+
+    def close(self):
+        if self.f is not None:
+            self.f.close()
+
+    def get_total_steps(self):
+        return self.total_steps
+
+    def get_episode_rewards(self):
+        return self.episode_rewards
+
+    def get_episode_lengths(self):
+        return self.episode_lengths
+
+    def get_episode_times(self):
+        return self.episode_times
+
+class LoadMonitorResultsError(Exception):
+    pass
+
+def get_monitor_files(dir):
+    return glob(osp.join(dir, "*" + Monitor.EXT))
+
+def load_monitor_log(dir):
+    import pandas
+    monitor_files = (
+        glob(osp.join(dir, "*monitor.json")) +
+        glob(osp.join(dir, "*monitor.csv"))) # get both csv and (old) json files
+    if not monitor_files:
+        raise LoadMonitorResultsError("no monitor files of the form *%s found in %s" % (Monitor.EXT, dir))
+    dfs = []
+    headers = []
+    for fname in monitor_files:
+        with open(fname, 'rt') as fh:
+            if fname.endswith('csv'):
+                firstline = fh.readline()
+                assert firstline[0] == '#'
+                header = json.loads(firstline[1:])
+                df = pandas.read_csv(fh, index_col=None)
+                headers.append(header)
+            elif fname.endswith('json'): # Deprecated json format
+                episodes = []
+                lines = fh.readlines()
+                header = json.loads(lines[0])
+                headers.append(header)
+                for line in lines[1:]:
+                    episode = json.loads(line)
+                    episodes.append(episode)
+                df = pandas.DataFrame(episodes)
+            else:
+                assert 0, 'unreachable'
+            df['t'] += header['t_start']
+        dfs.append(df)
+    df = pandas.concat(dfs)
+    df.sort_values('t', inplace=True)
+    df.reset_index(inplace=True)
+    df['t'] -= min(header['t_start'] for header in headers)
+    # df.headers = headers # HACK to preserve backwards compatibility
+    return df

deep_rl/component/policy.py

@@ -0,0 +1,57 @@
+#######################################################################
+# Copyright (C) 2017 Shangtong Zhang(zhangshangtong.cpp@gmail.com)    #
+# Permission given to modify the code as long as you keep this        #
+# declaration at the top                                              #
+#######################################################################
+
+import numpy as np
+
+class GreedyPolicy:
+    def __init__(self, epsilon, final_step, min_epsilon):
+        self.init_epsilon = self.epsilon = epsilon
+        self.current_steps = 0
+        self.min_epsilon = min_epsilon
+        self.final_step = final_step
+
+    def sample(self, action_value, deterministic=False):
+        if deterministic:
+            return np.argmax(action_value)
+        if np.random.rand() < self.epsilon:
+            return np.random.randint(0, len(action_value))
+        return np.argmax(action_value)
+
+    def update_epsilon(self):
+        self.epsilon = self.init_epsilon - float(self.current_steps) / self.final_step * (self.init_epsilon - self.min_epsilon)
+        self.epsilon = max(self.epsilon, self.min_epsilon)
+        self.current_steps += 1
+
+class StochasticGreedyPolicy:
+    def __init__(self, epsilons, final_step, min_epsilons, probs):
+        self.policies = []
+        self.probs = probs
+        for epsilon, min_epsilon in zip(epsilons, min_epsilons):
+            self.policies.append(GreedyPolicy(epsilon, final_step, min_epsilon))
+
+    def sample(self, action_value, deterministic=False):
+        return np.random.choice(self.policies, p=self.probs).sample(action_value, deterministic)
+
+    def update_epsilon(self):
+        for policy in self.policies:
+            policy.update_epsilon()
+
+class SamplePolicy:
+    def sample(self, action_value, deterministic=False):
+        if deterministic:
+            return np.argmax(action_value)
+        return np.random.choice(np.arange(len(action_value)), p=action_value)
+    def update_epsilon(self):
+        pass
+
+class GaussianPolicy:
+    def sample(self, mean, std, deterministic=False):
+        if deterministic:
+            return mean
+        return mean + std * np.random.randn(*mean.shape)
+
+    def update_epsilon(self):
+        pass

deep_rl/component/random_process.py

@@ -0,0 +1,13 @@
+import numpy as np
+
+class RandomProcess(object):
+    def reset_states(self):
+        pass
+
+class GaussianProcess(RandomProcess):
+    def __init__(self, size, std_schedule):
+        self.size = size
+        self.std_schedule = std_schedule
+
+    def sample(self):
+        return np.random.randn(self.size) * self.std_schedule()

deep_rl/component/replay.py

@@ -0,0 +1,85 @@
+#######################################################################
+# Copyright (C) 2017 Shangtong Zhang(zhangshangtong.cpp@gmail.com)    #
+# Permission given to modify the code as long as you keep this        #
+# declaration at the top                                              #
+#######################################################################
+
+import numpy as np
+
+class Replay:
+    def __init__(self, memory_size, batch_size):
+        self.memory_size = memory_size
+        self.batch_size = batch_size
+        self.data = None
+
+        self.pos = 0
+        self.full = False
+
+    def feed(self, experience):
+        if self.data is None:
+            self.data = []
+            for unit in experience:
+                if np.isscalar(unit):
+                    self.data.append(np.zeros(self.memory_size, dtype=type(unit)))
+                else:
+                    self.data.append(np.zeros((self.memory_size, ) + unit.shape, unit.dtype))
+        for buffer_unit, exp_unit in zip(self.data, experience):
+            buffer_unit[self.pos] = exp_unit
+
+        self.pos += 1
+        if self.pos == self.memory_size:
+            self.full = True
+            self.pos = 0
+
+    def feed_batch(self, experience):
+        experience = zip(*experience)
+        for exp in experience:
+            self.feed(exp)
+
+    def sample(self, batch_size=None):
+        if batch_size is None:
+            batch_size = self.batch_size
+        upper_bound = self.memory_size if self.full else self.pos
+        sampled_indices = np.random.randint(0, upper_bound, size=batch_size)
+        return [unit[sampled_indices] for unit in self.data]
+
+    def size(self):
+        if self.full:
+            return self.memory_size
+        return self.pos
+
+    def empty(self):
+        return not self.full and not self.pos
+
+class SkewedReplay:
+    def __init__(self, memory_size, batch_size):
+        memory_size = memory_size / 2
+        self.non_zero_reward = Replay(memory_size, batch_size / 2)
+        self.zero_reward = Replay(memory_size, batch_size / 2)
+        self.batch_size = batch_size
+
+    def feed(self, experiences):
+        experiences = zip(*experiences)
+        for exp in experiences:
+            if np.abs(exp[2]) < 1e-5:
+                self.zero_reward.feed(exp)
+            else:
+                self.non_zero_reward.feed(exp)
+
+    def sample(self):
+        if self.zero_reward.empty():
+            batch = self.non_zero_reward.sample(self.batch_size)
+        elif self.non_zero_reward.empty():
+            batch = self.zero_reward.sample(self.batch_size)
+        else:
+            non_zero_batch_size = min(self.non_zero_reward.size(), self.batch_size / 2)
+            zero_batch_size = min(self.zero_reward.size(), self.batch_size / 2)
+            batch1 = self.zero_reward.sample(zero_batch_size)
+            batch2 = self.non_zero_reward.sample(non_zero_batch_size)
+            batch = list(map(lambda seq: np.concatenate([np.asarray(x) for x in seq], axis=0), zip(batch1, batch2)))
+        batch = list(map(lambda x: np.asarray(x), batch))
+        return batch
+
+
+
+

deep_rl/component/task.py

@@ -0,0 +1,206 @@
+#######################################################################
+# Copyright (C) 2017 Shangtong Zhang(zhangshangtong.cpp@gmail.com)    #
+# Permission given to modify the code as long as you keep this        #
+# declaration at the top                                              #
+#######################################################################
+from .atari_wrapper import *
+import multiprocessing as mp
+import sys
+from .bench import Monitor
+from ..utils import *
+import uuid
+
+class BaseTask:
+    def set_monitor(self, env, log_dir):
+        if log_dir is None:
+            return env
+        mkdir(log_dir)
+        return Monitor(env, '%s/%s' % (log_dir, uuid.uuid4()))
+
+    def reset(self):
+        return self.env.reset()
+
+    def step(self, action):
+        return self.env.step(action)
+
+    def seed(self, random_seed):
+        return self.env.seed(random_seed)
+
+class ClassicalControl(BaseTask):
+    def __init__(self, name='CartPole-v0', max_steps=200, log_dir=None):
+        BaseTask.__init__(self)
+        self.name = name
+        self.env = gym.make(self.name)
+        self.env._max_episode_steps = max_steps
+        self.action_dim = self.env.action_space.n
+        self.state_dim = self.env.observation_space.shape[0]
+        self.env = self.set_monitor(self.env, log_dir)
+
+class PixelAtari(BaseTask):
+    def __init__(self, name, seed=0, log_dir=None,
+                 frame_skip=4, history_length=4, dataset=False):
+        BaseTask.__init__(self)
+        env = make_atari(name, frame_skip)
+        env.seed(seed)
+        if dataset:
+            env = DatasetEnv(env)
+            self.dataset_env = env
+        env = self.set_monitor(env, log_dir)
+        env = wrap_deepmind(env, history_length=history_length)
+        self.env = env
+        self.action_dim = self.env.action_space.n
+        self.state_dim = self.env.observation_space.shape
+        self.name = name
+
+class RamAtari(BaseTask):
+    def __init__(self, name, no_op, frame_skip, log_dir=None):
+        BaseTask.__init__(self)
+        self.name = name
+        env = gym.make(name)
+        assert 'NoFrameskip' in env.spec.id
+        env = self.set_monitor(env, log_dir)
+        env = EpisodicLifeEnv(env)
+        env = NoopResetEnv(env, noop_max=no_op)
+        env = SkipEnv(env, skip=frame_skip)
+        if 'FIRE' in env.unwrapped.get_action_meanings():
+            env = FireResetEnv(env)
+        self.env = env
+        self.action_dim = self.env.action_space.n
+        self.state_dim = 128
+
+class Pendulum(BaseTask):
+    def __init__(self, log_dir=None):
+        BaseTask.__init__(self)
+        self.name = 'Pendulum-v0'
+        self.env = gym.make(self.name)
+        self.action_dim = self.env.action_space.shape[0]
+        self.state_dim = self.env.observation_space.shape[0]
+        self.env = self.set_monitor(self.env, log_dir)
+
+    def step(self, action):
+        return BaseTask.step(self, np.clip(2 * action, -2, 2))
+
+class Box2DContinuous(BaseTask):
+    def __init__(self, name, log_dir=None):
+        BaseTask.__init__(self)
+        self.name = name
+        self.env = gym.make(self.name)
+        self.action_dim = self.env.action_space.shape[0]
+        self.state_dim = self.env.observation_space.shape[0]
+        self.env = self.set_monitor(self.env, log_dir)
+
+    def step(self, action):
+        return BaseTask.step(self, np.clip(action, -1, 1))
+
+class Roboschool(BaseTask):
+    def __init__(self, name, log_dir=None):
+        import roboschool
+        BaseTask.__init__(self)
+        self.name = name
+        self.env = gym.make(self.name)
+        self.action_dim = self.env.action_space.shape[0]
+        self.state_dim = self.env.observation_space.shape[0]
+        self.env = self.set_monitor(self.env, log_dir)
+
+    def step(self, action):
+        return BaseTask.step(self, np.clip(action, -1, 1))
+
+class DMControl(BaseTask):
+    def __init__(self, domain_name, task_name, log_dir=None):
+        from dm_control import suite
+        import dm_control2gym
+        BaseTask.__init__(self)
+
+        self.name = domain_name + '_' + task_name
+        self.env = dm_control2gym.make(domain_name, task_name)
+
+        self.action_dim = self.env.action_space.shape[0]
+        self.state_dim = self.env.observation_space.shape[0]
+        self.env = self.set_monitor(self.env, log_dir)
+
+class GymRobotics(BaseTask):
+    def __init__(self, name, log_dir=None):
+        BaseTask.__init__(self)
+
+        self.name = name
+        self.env = gym.make(name)
+
+        self.action_dim = self.env.action_space.shape[0]
+        self.state_dim = len(self.flatten_state(self.env.reset()))
+        self.env = self.set_monitor(self.env, log_dir)
+
+    def flatten_state(self, state):
+        flat = []
+        for key, value in state.items():
+            flat.append(state[key])
+        flat = np.concatenate(flat, axis=0)
+        return flat
+
+    def reset(self):
+        return self.flatten_state(self.env.reset())
+
+    def step(self, action):
+        next_state, reward, done, _ = self.env.step(action)
+        return self.flatten_state(next_state), reward, done, _
+
+def sub_task(parent_pipe, pipe, task_fn, rank, log_dir):
+    np.random.seed()
+    seed = np.random.randint(0, sys.maxsize)
+    parent_pipe.close()
+    task = task_fn(log_dir=log_dir)
+    task.seed(seed)
+    while True:
+        op, data = pipe.recv()
+        if op == 'step':
+            ob, reward, done, info = task.step(data)
+            if done:
+                ob = task.reset()
+            pipe.send([ob, reward, done, info])
+        elif op == 'reset':
+            pipe.send(task.reset())
+        elif op == 'exit':
+            pipe.close()
+            return
+        else:
+            assert False, 'Unknown Operation'
+
+class ParallelizedTask:
+    def __init__(self, task_fn, num_workers, log_dir=None):
+        self.task_fn = task_fn
+        self.task = task_fn(log_dir=None)
+        self.name = self.task.name
+        if log_dir is not None:
+            mkdir(log_dir)
+        self.pipes, worker_pipes = zip(*[mp.Pipe() for _ in range(num_workers)])
+        args = [(p, wp, task_fn, rank, log_dir)
+                for rank, (p, wp) in enumerate(zip(self.pipes, worker_pipes))]
+        self.workers = [mp.Process(target=sub_task, args=arg) for arg in args]
+        for p in self.workers: p.start()
+        for p in worker_pipes: p.close()
+        self.state_dim = self.task.state_dim
+        self.action_dim = self.task.action_dim
+
+    def step(self, actions):
+        for pipe, action in zip(self.pipes, actions):
+            pipe.send(('step', action))
+        results = [p.recv() for p in self.pipes]
+        results = map(lambda x: np.stack(x), zip(*results))
+        return results
+
+    def reset(self, i=None):
+        if i is None:
+            for pipe in self.pipes:
+                pipe.send(('reset', None))
+            results = [p.recv() for p in self.pipes]
+        else:
+            self.pipes[i].send(('reset', None))
+            results = self.pipes[i].recv()
+        return np.stack(results)
+
+    def close(self):
+        for pipe in self.pipes:
+            pipe.send(('exit', None))
+        for p in self.workers: p.join()
+
+    def normalize_state(self, state):
+        return self.task.normalize_state(state)