Upgrade DDPG

deep_rl/component/atari_wrapper.py

@@ -254,6 +254,22 @@ class DatasetEnv(gym.Wrapper):
         self.saved_obs.append(obs)
         return obs
 
+class RenderEnv(gym.Wrapper):
+    def __init__(self, env):
+        gym.Wrapper.__init__(self, env)
+        self.observation_space = spaces.Box(low=0, high=255,
+            shape=(self.env.unwrapped._render_height, self.env.unwrapped._render_width, 3), dtype=np.uint8)
+
+    def step(self, action):
+        _, reward, done, info = self.env.step(action)
+        obs = self.env.render('rgb_array')
+        return obs, reward, done, info
+
+    def reset(self):
+        self.env.reset()
+        obs = self.env.render('rgb_array')
+        return obs
+
 def make_atari(env_id, frame_skip=4):
     env = gym.make(env_id)
     assert 'NoFrameskip' in env.spec.id

deep_rl/component/task.py

@@ -118,6 +118,23 @@ class Bullet(BaseTask):
     def step(self, action):
         return BaseTask.step(self, np.clip(action, -1, 1))
 
+class PixelBullet(BaseTask):
+    def __init__(self, name, seed=0, log_dir=None, frame_skip=4, history_length=4):
+        import pybullet_envs
+        self.name = name
+        env = gym.make(name)
+        env.seed(seed)
+        env = RenderEnv(env)
+        env = self.set_monitor(env, log_dir)
+        env = SkipEnv(env, skip=frame_skip)
+        env = WarpFrame(env)
+        env = WrapPyTorch(env)
+        if history_length:
+            env = StackFrame(env, history_length)
+        self.action_dim = env.action_space.shape[0]
+        self.state_dim = env.observation_space.shape
+        self.env = env
+
 class ProcessTask:
     def __init__(self, task_fn, log_dir=None):
         self.pipe, worker_pipe = mp.Pipe()
@@ -171,8 +188,11 @@ class ProcessWrapper(mp.Process):
                 raise Exception('Unknown command')
 
 class ParallelizedTask:
-    def __init__(self, task_fn, num_workers, log_dir=None):
-        self.tasks = [ProcessTask(task_fn, log_dir) for _ in range(num_workers)]
+    def __init__(self, task_fn, num_workers, log_dir=None, single_process=False):
+        if single_process:
+            self.tasks = [task_fn(log_dir=log_dir) for _ in range(num_workers)]
+        else:
+            self.tasks = [ProcessTask(task_fn, log_dir) for _ in range(num_workers)]
         self.state_dim = self.tasks[0].state_dim
         self.action_dim = self.tasks[0].action_dim
         self.name = self.tasks[0].name
@@ -187,4 +207,4 @@ class ParallelizedTask:
         return np.stack(results)
 
     def close(self):
-        for task in self.tasks: task.close()
+        for task in self.tasks: task.close()