[rllib] Better support and add two-trainer example for multiagent (#2443)

This adds a simple DQN+PPO example for multi-agent. We don't do anything fancy here, just syncing weights between two separate trainers. This potentially is wasting some compute, but is very simple to set up.

It might be nice to share experience collection between the top-level trainers in the future.

python/ray/rllib/agents/agent.py

@@ -36,7 +36,10 @@ COMMON_CONFIG = {
     # Arguments to pass to the env creator
     "env_config": {},
     # Arguments to pass to model
-    "model": {},
+    "model": {
+        "use_lstm": False,
+        "max_seq_len": 20,
+    },
     # Arguments to pass to the rllib optimizer
     "optimizer": {},
     # Configure TF for single-process operation by default
@@ -57,8 +60,13 @@ COMMON_CONFIG = {
 
     # === Multiagent ===
     "multiagent": {
+        # Map from policy ids to tuples of (policy_graph_cls, obs_space,
+        # act_space, config). See policy_evaluator.py for more info.
         "policy_graphs": {},
+        # Function mapping agent ids to policy ids.
         "policy_mapping_fn": None,
+        # Optional whitelist of policies to train, or None for all policies.
+        "policies_to_train": None,
     },
 }
 
@@ -143,6 +151,7 @@ class Agent(Trainable):
             env_creator,
             self.config["multiagent"]["policy_graphs"] or policy_graph,
             policy_mapping_fn=self.config["multiagent"]["policy_mapping_fn"],
+            policies_to_train=self.config["multiagent"]["policies_to_train"],
             tf_session_creator=(session_creator
                                 if config["tf_session_args"] else None),
             batch_steps=config["sample_batch_size"],
@@ -239,6 +248,23 @@ class Agent(Trainable):
             lambda p: p.compute_single_action(obs, state, is_training=False)[0]
         )
 
+    def get_weights(self, policies=None):
+        """Return a dictionary of policy ids to weights.
+
+        Arguments:
+            policies (list): Optional list of policies to return weights for,
+                or None for all policies.
+        """
+        return self.local_evaluator.get_weights(policies)
+
+    def set_weights(self, weights):
+        """Set policy weights by policy id.
+
+        Arguments:
+            weights (dict): Map of policy ids to weights to set.
+        """
+        self.local_evaluator.set_weights(weights)
+
 
 class _MockAgent(Agent):
     """Mock agent for use in tests"""

python/ray/rllib/agents/dqn/dqn.py

@@ -166,7 +166,8 @@ class DQNAgent(Agent):
     def update_target_if_needed(self):
         if self.global_timestep - self.last_target_update_ts > \
                 self.config["target_network_update_freq"]:
-            self.local_evaluator.foreach_policy(lambda p, _: p.update_target())
+            self.local_evaluator.foreach_trainable_policy(
+                lambda p, _: p.update_target())
             self.last_target_update_ts = self.global_timestep
             self.num_target_updates += 1
 
@@ -179,11 +180,12 @@ class DQNAgent(Agent):
             self.update_target_if_needed()
 
         exp_vals = [self.exploration0.value(self.global_timestep)]
-        self.local_evaluator.foreach_policy(
+        self.local_evaluator.foreach_trainable_policy(
             lambda p, _: p.set_epsilon(exp_vals[0]))
         for i, e in enumerate(self.remote_evaluators):
             exp_val = self.explorations[i].value(self.global_timestep)
-            e.foreach_policy.remote(lambda p, _: p.set_epsilon(exp_val))
+            e.foreach_trainable_policy.remote(
+                lambda p, _: p.set_epsilon(exp_val))
             exp_vals.append(exp_val)
 
         if self.config["per_worker_exploration"]:

python/ray/rllib/agents/ppo/ppo.py

@@ -98,7 +98,14 @@ class PPOAgent(Agent):
     def _train(self):
         prev_steps = self.optimizer.num_steps_sampled
         fetches = self.optimizer.step()
-        self.local_evaluator.for_policy(lambda pi: pi.update_kl(fetches["kl"]))
+        if "kl" in fetches:
+            # single-agent
+            self.local_evaluator.for_policy(
+                lambda pi: pi.update_kl(fetches["kl"]))
+        else:
+            # multi-agent
+            self.local_evaluator.foreach_trainable_policy(
+                lambda pi, pi_id: pi.update_kl(fetches[pi_id]["kl"]))
         FilterManager.synchronize(self.local_evaluator.filters,
                                   self.remote_evaluators)
         res = self.optimizer.collect_metrics()

python/ray/rllib/agents/ppo/ppo_policy_graph.py

@@ -4,6 +4,7 @@ from __future__ import print_function
 
 import tensorflow as tf
 
+import ray
 from ray.rllib.evaluation.postprocessing import compute_advantages
 from ray.rllib.evaluation.tf_policy_graph import TFPolicyGraph
 from ray.rllib.models.catalog import ModelCatalog
@@ -96,6 +97,7 @@ class PPOPolicyGraph(TFPolicyGraph):
             existing_inputs (list): Optional list of tuples that specify the
                 placeholders upon which the graph should be built upon.
         """
+        config = dict(ray.rllib.agents.ppo.ppo.DEFAULT_CONFIG, **config)
         self.sess = tf.get_default_session()
         self.action_space = action_space
         self.config = config

python/ray/rllib/evaluation/policy_evaluator.py

@@ -86,6 +86,7 @@ class PolicyEvaluator(EvaluatorInterface):
                  env_creator,
                  policy_graph,
                  policy_mapping_fn=None,
+                 policies_to_train=None,
                  tf_session_creator=None,
                  batch_steps=100,
                  batch_mode="truncate_episodes",
@@ -113,6 +114,8 @@ class PolicyEvaluator(EvaluatorInterface):
                 policy ids in multi-agent mode. This function will be called
                 each time a new agent appears in an episode, to bind that agent
                 to a policy for the duration of the episode.
+            policies_to_train (list): Optional whitelist of policies to train,
+                or None for all policies.
             tf_session_creator (func): A function that returns a TF session.
                 This is optional and only useful with TFPolicyGraph.
             batch_steps (int): The target number of env transitions to include
@@ -159,7 +162,6 @@ class PolicyEvaluator(EvaluatorInterface):
         policy_mapping_fn = (policy_mapping_fn
                              or (lambda agent_id: DEFAULT_POLICY_ID))
         self.env_creator = env_creator
-        self.policy_graph = policy_graph
         self.batch_steps = batch_steps
         self.batch_mode = batch_mode
         self.compress_observations = compress_observations
@@ -191,6 +193,7 @@ class PolicyEvaluator(EvaluatorInterface):
 
         self.tf_sess = None
         policy_dict = _validate_and_canonicalize(policy_graph, self.env)
+        self.policies_to_train = policies_to_train or list(policy_dict.keys())
         if _has_tensorflow_graph(policy_dict):
             with tf.Graph().as_default():
                 if tf_session_creator:
@@ -300,6 +303,16 @@ class PolicyEvaluator(EvaluatorInterface):
 
         return [func(policy, pid) for pid, policy in self.policy_map.items()]
 
+    def foreach_trainable_policy(self, func):
+        """Apply the given function to each (policy, policy_id) tuple.
+
+        This only applies func to policies in `self.policies_to_train`."""
+
+        return [
+            func(policy, pid) for pid, policy in self.policy_map.items()
+            if pid in self.policies_to_train
+        ]
+
     def sync_filters(self, new_filters):
         """Changes self's filter to given and rebases any accumulated delta.
 
@@ -326,10 +339,12 @@ class PolicyEvaluator(EvaluatorInterface):
                 f.clear_buffer()
         return return_filters
 
-    def get_weights(self):
+    def get_weights(self, policies=None):
+        if policies is None:
+            policies = self.policy_map.keys()
         return {
             pid: policy.get_weights()
-            for pid, policy in self.policy_map.items()
+            for pid, policy in self.policy_map.items() if pid in policies
         }
 
     def set_weights(self, weights):
@@ -342,6 +357,8 @@ class PolicyEvaluator(EvaluatorInterface):
             if self.tf_sess is not None:
                 builder = TFRunBuilder(self.tf_sess, "compute_gradients")
                 for pid, batch in samples.policy_batches.items():
+                    if pid not in self.policies_to_train:
+                        continue
                     grad_out[pid], info_out[pid] = (
                         self.policy_map[pid].build_compute_gradients(
                             builder, batch))
@@ -381,6 +398,8 @@ class PolicyEvaluator(EvaluatorInterface):
             if self.tf_sess is not None:
                 builder = TFRunBuilder(self.tf_sess, "compute_apply")
                 for pid, batch in samples.policy_batches.items():
+                    if pid not in self.policies_to_train:
+                        continue
                     info_out[pid], _ = (
                         self.policy_map[pid].build_compute_apply(
                             builder, batch))

python/ray/rllib/examples/multiagent_two_trainers.py

@@ -0,0 +1,101 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+"""Example of using two different training methods at once in multi-agent.
+
+Here we create a number of CartPole agents, some of which are trained with
+DQN, and some of which are trained with PPO. We periodically sync weights
+between the two trainers (note that no such syncing is needed when using just
+a single training method).
+
+For a simpler example, see also: multiagent_cartpole.py
+"""
+
+import argparse
+import gym
+
+import ray
+from ray.rllib.agents.dqn.dqn import DQNAgent
+from ray.rllib.agents.dqn.dqn_policy_graph import DQNPolicyGraph
+from ray.rllib.agents.ppo.ppo import PPOAgent
+from ray.rllib.agents.ppo.ppo_policy_graph import PPOPolicyGraph
+from ray.rllib.test.test_multi_agent_env import MultiCartpole
+from ray.tune.logger import pretty_print
+from ray.tune.registry import register_env
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--num-iters", type=int, default=20)
+
+if __name__ == "__main__":
+    args = parser.parse_args()
+    ray.init()
+
+    # Simple environment with 4 independent cartpole entities
+    register_env("multi_cartpole", lambda _: MultiCartpole(4))
+    single_env = gym.make("CartPole-v0")
+    obs_space = single_env.observation_space
+    act_space = single_env.action_space
+
+    # You can also have multiple policy graphs per trainer, but here we just
+    # show one each for PPO and DQN.
+    policy_graphs = {
+        "ppo_policy": (PPOPolicyGraph, obs_space, act_space, {}),
+        "dqn_policy": (DQNPolicyGraph, obs_space, act_space, {}),
+    }
+
+    def policy_mapping_fn(agent_id):
+        if agent_id % 2 == 0:
+            return "ppo_policy"
+        else:
+            return "dqn_policy"
+
+    ppo_trainer = PPOAgent(
+        env="multi_cartpole",
+        config={
+            "multiagent": {
+                "policy_graphs": policy_graphs,
+                "policy_mapping_fn": policy_mapping_fn,
+                "policies_to_train": ["ppo_policy"],
+            },
+            "simple_optimizer": True,
+            # disable filters, otherwise we would need to synchronize those
+            # as well to the DQN agent
+            "observation_filter": "NoFilter",
+        })
+
+    dqn_trainer = DQNAgent(
+        env="multi_cartpole",
+        config={
+            "multiagent": {
+                "policy_graphs": policy_graphs,
+                "policy_mapping_fn": policy_mapping_fn,
+                "policies_to_train": ["dqn_policy"],
+            },
+            "gamma": 0.95,
+            "n_step": 3,
+        })
+
+    # disable DQN exploration when used by the PPO trainer
+    ppo_trainer.optimizer.foreach_evaluator(
+        lambda ev: ev.for_policy(
+            lambda pi: pi.set_epsilon(0.0), policy_id="dqn_policy"))
+
+    # You should see both the printed X and Y approach 200 as this trains:
+    # info:
+    #   policy_reward_mean:
+    #     dqn_policy: X
+    #     ppo_policy: Y
+    for i in range(args.num_iters):
+        print("== Iteration", i, "==")
+
+        # improve the DQN policy
+        print("-- DQN --")
+        print(pretty_print(dqn_trainer.train()))
+
+        # improve the PPO policy
+        print("-- PPO --")
+        print(pretty_print(ppo_trainer.train()))
+
+        # swap weights to synchronize
+        dqn_trainer.set_weights(ppo_trainer.get_weights(["ppo_policy"]))
+        ppo_trainer.set_weights(dqn_trainer.get_weights(["dqn_policy"]))

python/ray/rllib/optimizers/multi_gpu_optimizer.py

@@ -58,13 +58,15 @@ class LocalMultiGPUOptimizer(PolicyOptimizer):
 
         print("LocalMultiGPUOptimizer devices", self.devices)
 
-        assert set(self.local_evaluator.policy_map.keys()) == {"default"}, \
-            ("Multi-agent is not supported with multi-GPU. Try using the "
-             "simple optimizer instead.")
+        if set(self.local_evaluator.policy_map.keys()) != {"default"}:
+            raise ValueError(
+                "Multi-agent is not supported with multi-GPU. Try using the "
+                "simple optimizer instead.")
         self.policy = self.local_evaluator.policy_map["default"]
-        assert isinstance(self.policy, TFPolicyGraph), \
-            ("Only TF policies are supported with multi-GPU. Try using the "
-             "simple optimizer instead.")
+        if not isinstance(self.policy, TFPolicyGraph):
+            raise ValueError(
+                "Only TF policies are supported with multi-GPU. Try using the "
+                "simple optimizer instead.")
 
         # per-GPU graph copies created below must share vars with the policy
         # reuse is set to AUTO_REUSE because Adam nodes are created after