From 9d53c75bd8b1e668812f05a2009ab7cddbb3b852 Mon Sep 17 00:00:00 2001 From: Shangtong Zhang Date: Thu, 10 May 2018 22:30:31 -0600 Subject: [PATCH] Replace Mujoco with PyBullet --- README.md | 4 ++- deep_rl/agent/BaseAgent.py | 2 ++ deep_rl/component/task.py | 66 +++++++++++++++++++++++--------------- examples.py | 21 ++++-------- 4 files changed, 53 insertions(+), 40 deletions(-) diff --git a/README.md b/README.md index 8021b6d..2b98a6c 100644 --- a/README.md +++ b/README.md @@ -28,12 +28,14 @@ Support for PyTorch v0.3.x can be found in [v0.2](https://github.com/ShangtongZh * PyTorch v0.4.0 * Python 3.6, 3.5 or 2.7 (deprecated) * Core dependencies: `pip install -e .` -* Optional: [Roboschool](https://github.com/openai/roboschool), [DeepMind Control Suite](https://github.com/deepmind/dm_control)+[DMControl2Gym](dm_control2gym) +* Optional: [Roboschool](https://github.com/openai/roboschool), [PyBullet](https://pypi.org/project/pybullet/) # Usage ```examples.py``` contains examples for all the implemented algorithms +```Dockerfile``` contains an example environment (w/ pybullet, w/o roboschool, w/o GPU) + Please use this bibtex if you want to cite this repo ``` @misc{deeprl, diff --git a/deep_rl/agent/BaseAgent.py b/deep_rl/agent/BaseAgent.py index a412839..d5e8b07 100644 --- a/deep_rl/agent/BaseAgent.py +++ b/deep_rl/agent/BaseAgent.py @@ -18,6 +18,8 @@ class BaseAgent: def close(self): if hasattr(self.task, 'close'): self.task.close() + if hasattr(self.evaluation_env, 'close'): + self.evaluation_env.close() def save(self, filename): torch.save(self.network.state_dict(), filename) diff --git a/deep_rl/component/task.py b/deep_rl/component/task.py index 6cbcb99..c3d1803 100644 --- a/deep_rl/component/task.py +++ b/deep_rl/component/task.py @@ -105,12 +105,44 @@ class Roboschool(BaseTask): def step(self, action): return BaseTask.step(self, np.clip(action, -1, 1)) +class Bullet(BaseTask): + def __init__(self, name, log_dir=None): + import pybullet_envs + BaseTask.__init__(self) + self.name = name + self.env = gym.make(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 ProcessTask: + def __init__(self, task_fn, log_dir): + self.pipe, worker_pipe = mp.Pipe() + self.worker = ProcessWrapper(worker_pipe, task_fn, log_dir) + self.worker.start() + self.pipe.send([ProcessWrapper.SPECS, None]) + self.state_dim, self.action_dim, self.name = self.pipe.recv() + + def step(self, action): + self.pipe.send([ProcessWrapper.STEP, action]) + return self.pipe.recv() + + def reset(self): + self.pipe.send([ProcessWrapper.RESET, None]) + return self.pipe.recv() + + def close(self): + self.pipe.send([ProcessWrapper.EXIT, None]) + class ProcessWrapper(mp.Process): STEP = 0 RESET = 1 EXIT = 2 SPECS = 3 - def __init__(self, pipe, task_fn, rank, log_dir): + def __init__(self, pipe, task_fn, log_dir): mp.Process.__init__(self) self.pipe = pipe self.task_fn = task_fn @@ -140,35 +172,19 @@ class ProcessWrapper(mp.Process): class ParallelizedTask: def __init__(self, task_fn, num_workers, log_dir=None): - self.task_fn = task_fn - if log_dir is not None: - mkdir(log_dir) - self.pipes, worker_pipes = zip(*[mp.Pipe() for _ in range(num_workers)]) - args = [(wp, task_fn, rank, log_dir) - for rank, wp in enumerate(worker_pipes)] - self.workers = [ProcessWrapper(*arg) for arg in args] - for p in self.workers: p.start() - self.pipes[0].send([ProcessWrapper.SPECS, None]) - self.state_dim, self.action_dim, self.name = self.pipes[0].recv() + 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 def step(self, actions): - for pipe, action in zip(self.pipes, actions): - pipe.send((ProcessWrapper.STEP, action)) - results = [p.recv() for p in self.pipes] + results = [task.step(action) for task, action in zip(self.tasks, actions)] 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((ProcessWrapper.RESET, None)) - results = [p.recv() for p in self.pipes] - else: - self.pipes[i].send((ProcessWrapper.RESET, None)) - results = self.pipes[i].recv() + def reset(self): + results = [task.reset() for task in self.tasks] return np.stack(results) def close(self): - for pipe in self.pipes: - pipe.send((ProcessWrapper.EXIT, None)) - for p in self.workers: p.join() + for task in self.tasks: task.close() diff --git a/examples.py b/examples.py index 331fb2b..5c6fd16 100644 --- a/examples.py +++ b/examples.py @@ -312,12 +312,7 @@ def ppo_continuous(): config = Config() config.num_workers = 1 # task_fn = lambda log_dir: Pendulum(log_dir=log_dir) - # task_fn = lambda log_dir: Roboschool('RoboschoolInvertedPendulum-v1', log_dir=log_dir) - task_fn = lambda log_dir: Roboschool('RoboschoolAnt-v1', log_dir=log_dir) - # task_fn = lambda log_dir: Roboschool('RoboschoolReacher-v1', log_dir=log_dir) - # task_fn = lambda log_dir: Roboschool('RoboschoolHopper-v1', log_dir=log_dir) - # task_fn = lambda log_dir: DMControl('cartpole', 'balance', log_dir=log_dir) - # task_fn = lambda log_dir: DMControl('hopper', 'hop', log_dir=log_dir) + task_fn = lambda log_dir: Bullet('AntBulletEnv-v0', log_dir=log_dir) config.task_fn = lambda: ParallelizedTask(task_fn, config.num_workers, log_dir=get_default_log_dir(ppo_continuous.__name__)) actor_network_fn = lambda state_dim, action_dim: GaussianActorNet( action_dim, FCBody(state_dim)) @@ -345,14 +340,12 @@ def ddpg_continuous(): config = Config() log_dir = get_default_log_dir(ddpg_continuous.__name__) # config.task_fn = lambda: Pendulum(log_dir=log_dir) - # config.task_fn = lambda: Roboschool('RoboschoolInvertedPendulum-v1', log_dir=log_dir) - # config.task_fn = lambda: Roboschool('RoboschoolReacher-v1', log_dir=log_dir) - config.task_fn = lambda: Roboschool('RoboschoolHopper-v1') - # config.task_fn = lambda: Roboschool('RoboschoolAnt-v1', log_dir=log_dir) - # config.task_fn = lambda: Roboschool('RoboschoolWalker2d-v1', log_dir=log_dir) - # config.task_fn = lambda: DMControl('cartpole', 'balance', log_dir=log_dir) - # config.task_fn = lambda: DMControl('finger', 'spin', log_dir=log_dir) - config.evaluation_env = Roboschool('RoboschoolHopper-v1', log_dir=log_dir) + config.task_fn = lambda **kwargs: Bullet('AntBulletEnv-v0', **kwargs) + + # start the test environment in a new process, it is a workaround to the issue + # https://github.com/bulletphysics/bullet3/issues/1643 + config.evaluation_env = ProcessTask(config.task_fn, log_dir=log_dir) + config.actor_network_fn = lambda state_dim, action_dim: DeterministicActorNet( action_dim, FCBody(state_dim, (300, 200))) config.critic_network_fn = lambda state_dim, action_dim: DeterministicCriticNet(