Upgrade DDPG

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