From 4af0aa62582f3e7eed18c1778277fe1506f1be5a Mon Sep 17 00:00:00 2001 From: Philipp Moritz Date: Tue, 14 Mar 2017 13:31:29 -0700 Subject: [PATCH] Atari on pixels (#364) * pong on pixels working (not cleaned up) * make training compatible with all atari games * cartpole runs * Update documentation and usage for policy gradients. --- doc/source/example-policy-gradient.rst | 7 +- examples/policy_gradient/examples/example.py | 110 ++++++++++++------- examples/policy_gradient/reinforce/agent.py | 10 +- examples/policy_gradient/reinforce/env.py | 30 ++++- examples/policy_gradient/reinforce/policy.py | 10 +- 5 files changed, 110 insertions(+), 57 deletions(-) diff --git a/doc/source/example-policy-gradient.rst b/doc/source/example-policy-gradient.rst index 449a83848..92d049bcc 100644 --- a/doc/source/example-policy-gradient.rst +++ b/doc/source/example-policy-gradient.rst @@ -23,9 +23,12 @@ Then you can run the example as follows. .. code-block:: bash - python ray/examples/policy_gradient/examples/example.py + python ray/examples/policy_gradient/examples/example.py --environment=Pong-ram-v3 -This will train an agent on an Atari environment. +This will train an agent on the ``Pong-ram-v3`` Atari environment. You can also +try passing in the ``Pong-v0`` environment or the ``CartPole-v0`` environment. +If you wish to use a different environment, you will need to change a few lines +in ``example.py``. .. _`TensorFlow with GPU support`: https://www.tensorflow.org/install/ .. _`code for this example`: https://github.com/ray-project/ray/tree/master/examples/policy_gradient diff --git a/examples/policy_gradient/examples/example.py b/examples/policy_gradient/examples/example.py index 833ceb43b..49f823206 100644 --- a/examples/policy_gradient/examples/example.py +++ b/examples/policy_gradient/examples/example.py @@ -2,8 +2,10 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import argparse import ray +from reinforce.env import NoPreprocessor, AtariRamPreprocessor, AtariPixelPreprocessor from reinforce.agent import Agent, RemoteAgent from reinforce.rollout import collect_samples from reinforce.utils import iterate, shuffle @@ -17,49 +19,73 @@ config = {"kl_coeff": 0.2, "kl_target": 0.01, "timesteps_per_batch": 40000} -ray.init() +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Run the policy gradient algorithm.") + parser.add_argument("--environment", default="Pong-v0", type=str, + help="The gym environment to use.") + parser.add_argument("--redis-address", default=None, type=str, + help="The Redis address of the cluster.") -mdp_name = "Pong-ram-v3" + args = parser.parse_args() -agents = [RemoteAgent(mdp_name, 1, config, False) for _ in range(5)] -agent = Agent(mdp_name, 1, config, True) + ray.init(redis_address=args.redis_address) -kl_coeff = config["kl_coeff"] + ray.register_class(AtariRamPreprocessor) + ray.register_class(AtariPixelPreprocessor) + ray.register_class(NoPreprocessor) -for j in range(1000): - print("== iteration", j) - weights = agent.get_weights() - [a.load_weights(weights) for a in agents] - trajectory, total_reward, traj_len_mean = collect_samples(agents, config["timesteps_per_batch"], 0.995, 1.0, 2000) - print("total reward is ", total_reward) - print("trajectory length mean is ", traj_len_mean) - print("timesteps: ", trajectory["dones"].shape[0]) - trajectory["advantages"] = (trajectory["advantages"] - trajectory["advantages"].mean()) / trajectory["advantages"].std() - print("Computing policy (optimizer='" + agent.optimizer.get_name() + "', iterations=" + str(config["num_sgd_iter"]) + ", stepsize=" + str(config["sgd_stepsize"]) + "):") - names = ["iter", "loss", "kl", "entropy"] - print(("{:>15}" * len(names)).format(*names)) - trajectory = shuffle(trajectory) - ppo = agent.ppo - for i in range(config["num_sgd_iter"]): - # Test on current set of rollouts - loss, kl, entropy = agent.sess.run([ppo.loss, ppo.mean_kl, ppo.mean_entropy], - feed_dict={ppo.observations: trajectory["observations"], - ppo.advantages: trajectory["advantages"], - ppo.actions: trajectory["actions"].squeeze(), - ppo.prev_logits: trajectory["logprobs"], - ppo.kl_coeff: kl_coeff}) - print("{:>15}{:15.5e}{:15.5e}{:15.5e}".format(i, loss, kl, entropy)) - # Run SGD for training on current set of rollouts - for batch in iterate(trajectory, config["sgd_batchsize"]): - agent.sess.run([agent.train_op], - feed_dict={ppo.observations: batch["observations"], - ppo.advantages: batch["advantages"], - ppo.actions: batch["actions"].squeeze(), - ppo.prev_logits: batch["logprobs"], - ppo.kl_coeff: kl_coeff}) - if kl > 2.0 * config["kl_target"]: - kl_coeff *= 1.5 - elif kl < 0.5 * config["kl_target"]: - kl_coeff *= 0.5 - print("kl div = ", kl) - print("kl coeff = ", kl_coeff) + mdp_name = args.environment + if args.environment == "Pong-v0": + preprocessor = AtariPixelPreprocessor() + elif mdp_name == "Pong-ram-v3": + preprocessor = AtariRamPreprocessor() + elif mdp_name == "CartPole-v0": + preprocessor = NoPreprocessor() + else: + print("No environment was chosen, so defaulting to Pong-v0.") + mdp_name = "Pong-v0" + preprocessor = AtariPixelPreprocessor() + + print("Using the environment {}.".format(mdp_name)) + agents = [RemoteAgent(mdp_name, 1, preprocessor, config, False) for _ in range(5)] + agent = Agent(mdp_name, 1, preprocessor, config, True) + + kl_coeff = config["kl_coeff"] + + for j in range(1000): + print("== iteration", j) + weights = ray.put(agent.get_weights()) + [a.load_weights(weights) for a in agents] + trajectory, total_reward, traj_len_mean = collect_samples(agents, config["timesteps_per_batch"], 0.995, 1.0, 2000) + print("total reward is ", total_reward) + print("trajectory length mean is ", traj_len_mean) + print("timesteps: ", trajectory["dones"].shape[0]) + trajectory["advantages"] = (trajectory["advantages"] - trajectory["advantages"].mean()) / trajectory["advantages"].std() + print("Computing policy (optimizer='" + agent.optimizer.get_name() + "', iterations=" + str(config["num_sgd_iter"]) + ", stepsize=" + str(config["sgd_stepsize"]) + "):") + names = ["iter", "loss", "kl", "entropy"] + print(("{:>15}" * len(names)).format(*names)) + trajectory = shuffle(trajectory) + ppo = agent.ppo + for i in range(config["num_sgd_iter"]): + # Test on current set of rollouts + loss, kl, entropy = agent.sess.run([ppo.loss, ppo.mean_kl, ppo.mean_entropy], + feed_dict={ppo.observations: trajectory["observations"], + ppo.advantages: trajectory["advantages"], + ppo.actions: trajectory["actions"].squeeze(), + ppo.prev_logits: trajectory["logprobs"], + ppo.kl_coeff: kl_coeff}) + print("{:>15}{:15.5e}{:15.5e}{:15.5e}".format(i, loss, kl, entropy)) + # Run SGD for training on current set of rollouts + for batch in iterate(trajectory, config["sgd_batchsize"]): + agent.sess.run([agent.train_op], + feed_dict={ppo.observations: batch["observations"], + ppo.advantages: batch["advantages"], + ppo.actions: batch["actions"].squeeze(), + ppo.prev_logits: batch["logprobs"], + ppo.kl_coeff: kl_coeff}) + if kl > 2.0 * config["kl_target"]: + kl_coeff *= 1.5 + elif kl < 0.5 * config["kl_target"]: + kl_coeff *= 0.5 + print("kl div = ", kl) + print("kl coeff = ", kl_coeff) diff --git a/examples/policy_gradient/reinforce/agent.py b/examples/policy_gradient/reinforce/agent.py index 82f313c6d..c023bcada 100644 --- a/examples/policy_gradient/reinforce/agent.py +++ b/examples/policy_gradient/reinforce/agent.py @@ -14,16 +14,18 @@ from reinforce.rollout import rollouts, add_advantage_values class Agent(object): - def __init__(self, name, batchsize, config, use_gpu): + def __init__(self, name, batchsize, preprocessor, config, use_gpu): if not use_gpu: os.environ["CUDA_VISIBLE_DEVICES"] = "" - self.env = BatchedEnv(name, batchsize, preprocessor=None) + self.env = BatchedEnv(name, batchsize, preprocessor=preprocessor) + if preprocessor.shape is None: + preprocessor.shape = self.env.observation_space.shape self.sess = tf.Session() - self.ppo = ProximalPolicyLoss(self.env.observation_space, self.env.action_space, config, self.sess) + self.ppo = ProximalPolicyLoss(self.env.observation_space, self.env.action_space, preprocessor, config, self.sess) self.optimizer = tf.train.AdamOptimizer(config["sgd_stepsize"]) self.train_op = self.optimizer.minimize(self.ppo.loss) self.variables = ray.experimental.TensorFlowVariables(self.ppo.loss, self.sess) - self.observation_filter = MeanStdFilter(self.env.observation_space.shape, clip=None) + self.observation_filter = MeanStdFilter(preprocessor.shape, clip=None) self.reward_filter = MeanStdFilter((), clip=5.0) self.sess.run(tf.global_variables_initializer()) diff --git a/examples/policy_gradient/reinforce/env.py b/examples/policy_gradient/reinforce/env.py index 30649c029..2acd4f8f7 100644 --- a/examples/policy_gradient/reinforce/env.py +++ b/examples/policy_gradient/reinforce/env.py @@ -5,12 +5,30 @@ from __future__ import print_function import gym import numpy as np -def atari_preprocessor(observation): - "Convert images from (210, 160, 3) to (3, 80, 80) by downsampling." - return (observation[25:-25:2,::2,:][None] - 128.0) / 128.8 +class AtariPixelPreprocessor(object): -def ram_preprocessor(observation): - return (observation - 128.0) / 128.0 + def __init__(self): + self.shape = (80, 80, 3) + + def __call__(self, observation): + "Convert images from (210, 160, 3) to (3, 80, 80) by downsampling." + return (observation[25:-25:2,::2,:][None] - 128.0) / 128.8 + +class AtariRamPreprocessor(object): + + def __init__(self): + self.shape = (128,) + + def __call__(self, observation): + return (observation - 128.0) / 128.0 + +class NoPreprocessor(object): + + def __init__(self): + self.shape = None + + def __call__(self, observation): + return observation class BatchedEnv(object): "A BatchedEnv holds multiple gym enviroments and performs steps on all of them." @@ -36,7 +54,7 @@ class BatchedEnv(object): observations.append(np.zeros(self.shape)) rewards.append(0.0) continue - observation, reward, done, info = self.envs[i].step(action) + observation, reward, done, info = self.envs[i].step(action if len(action) > 1 else action[0]) if render: self.envs[0].render() observations.append(self.preprocessor(observation)) diff --git a/examples/policy_gradient/reinforce/policy.py b/examples/policy_gradient/reinforce/policy.py index af0d613e9..11ce7f8e3 100644 --- a/examples/policy_gradient/reinforce/policy.py +++ b/examples/policy_gradient/reinforce/policy.py @@ -11,11 +11,11 @@ from reinforce.distributions import Categorical, DiagGaussian class ProximalPolicyLoss(object): - def __init__(self, observation_space, action_space, config, sess): + def __init__(self, observation_space, action_space, preprocessor, config, sess): assert isinstance(action_space, gym.spaces.Discrete) or isinstance(action_space, gym.spaces.Box) # adapting the kl divergence self.kl_coeff = tf.placeholder(name="newkl", shape=(), dtype=tf.float32) - self.observations = tf.placeholder(tf.float32, shape=(None,) + observation_space.shape) + self.observations = tf.placeholder(tf.float32, shape=(None,) + preprocessor.shape) self.advantages = tf.placeholder(tf.float32, shape=(None,)) if isinstance(action_space, gym.spaces.Box): @@ -33,7 +33,11 @@ class ProximalPolicyLoss(object): raise NotImplemented("action space" + str(type(env.action_space)) + "currently not supported") self.prev_logits = tf.placeholder(tf.float32, shape=(None, self.logit_dim)) self.prev_dist = Distribution(self.prev_logits) - self.curr_logits = fc_net(self.observations, num_classes=self.logit_dim) + if len(observation_space.shape) > 1: + self.curr_logits = vision_net(self.observations, num_classes=self.logit_dim) + else: + assert len(observation_space.shape) == 1 + self.curr_logits = fc_net(self.observations, num_classes=self.logit_dim) self.curr_dist = Distribution(self.curr_logits) self.sampler = self.curr_dist.sample() self.entropy = self.curr_dist.entropy()