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

@@ -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

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)

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

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

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