[rllib] Add custom value functions, fix up and document multi-agent variable sharing (#3151)

python/ray/rllib/examples/multiagent_cartpole.py

@@ -16,9 +16,13 @@ import argparse
 import gym
 import random
 
+import tensorflow as tf
+import tensorflow.contrib.slim as slim
+
 import ray
 from ray import tune
-from ray.rllib.agents.pg.pg_policy_graph import PGPolicyGraph
+from ray.rllib.agents.ppo.ppo_policy_graph import PPOPolicyGraph
+from ray.rllib.models import Model, ModelCatalog
 from ray.rllib.test.test_multi_agent_env import MultiCartpole
 from ray.tune import run_experiments
 from ray.tune.registry import register_env
@@ -29,26 +33,65 @@ parser.add_argument("--num-agents", type=int, default=4)
 parser.add_argument("--num-policies", type=int, default=2)
 parser.add_argument("--num-iters", type=int, default=20)
 
+
+class CustomModel1(Model):
+    def _build_layers_v2(self, input_dict, num_outputs, options):
+        # Example of (optional) weight sharing between two different policies.
+        # Here, we share the variables defined in the 'shared' variable scope
+        # by entering it explicitly with tf.AUTO_REUSE. This creates the
+        # variables for the 'fc1' layer in a global scope called 'shared'
+        # outside of the policy's normal variable scope.
+        with tf.variable_scope(
+                tf.VariableScope(tf.AUTO_REUSE, "shared"),
+                reuse=tf.AUTO_REUSE,
+                auxiliary_name_scope=False):
+            last_layer = slim.fully_connected(
+                input_dict["obs"], 64, activation_fn=tf.nn.relu, scope="fc1")
+        output = slim.fully_connected(
+            last_layer, num_outputs, activation_fn=None, scope="fc_out")
+        return output, last_layer
+
+
+class CustomModel2(Model):
+    def _build_layers_v2(self, input_dict, num_outputs, options):
+        # Weights shared with CustomModel1
+        with tf.variable_scope(
+                tf.VariableScope(tf.AUTO_REUSE, "shared"),
+                reuse=tf.AUTO_REUSE,
+                auxiliary_name_scope=False):
+            last_layer = slim.fully_connected(
+                input_dict["obs"], 64, activation_fn=tf.nn.relu, scope="fc1")
+        output = slim.fully_connected(
+            last_layer, num_outputs, activation_fn=None, scope="fc_out")
+        return output, last_layer
+
+
 if __name__ == "__main__":
     args = parser.parse_args()
     ray.init()
 
     # Simple environment with `num_agents` independent cartpole entities
     register_env("multi_cartpole", lambda _: MultiCartpole(args.num_agents))
+    ModelCatalog.register_custom_model("model1", CustomModel1)
+    ModelCatalog.register_custom_model("model2", CustomModel2)
     single_env = gym.make("CartPole-v0")
     obs_space = single_env.observation_space
     act_space = single_env.action_space
 
-    def gen_policy():
+    # Each policy can have a different configuration (including custom model)
+    def gen_policy(i):
         config = {
+            "model": {
+                "custom_model": ["model1", "model2"][i % 2],
+            },
             "gamma": random.choice([0.5, 0.8, 0.9, 0.95, 0.99]),
             "n_step": random.choice([1, 2, 3, 4, 5]),
         }
-        return (PGPolicyGraph, obs_space, act_space, config)
+        return (PPOPolicyGraph, obs_space, act_space, config)
 
-    # Setup PG with an ensemble of `num_policies` different policy graphs
+    # Setup PPO with an ensemble of `num_policies` different policy graphs
     policy_graphs = {
-        "policy_{}".format(i): gen_policy()
+        "policy_{}".format(i): gen_policy(i)
         for i in range(args.num_policies)
     }
     policy_ids = list(policy_graphs.keys())