[rllib] Add test for multi-agent support and fix IMPALA multi-agent (#3289)

IMPALA support for multiagent was broken since IMPALA has a requirement that batch sizes be of a certain length. However multi-agent envs can create variable-length batches.

Fix this by adding zero-padding as needed (similar to the RNN case).

python/ray/rllib/test/test_supported_spaces.py

@@ -9,6 +9,7 @@ import sys
 
 import ray
 from ray.rllib.agents.agent import get_agent_class
+from ray.rllib.test.test_multi_agent_env import MultiCartpole, MultiMountainCar
 from ray.rllib.utils.error import UnsupportedSpaceException
 from ray.tune.registry import register_env
 
@@ -88,9 +89,27 @@ def check_support(alg, config, stats, check_bounds=False):
             stats[alg, a_name, o_name] = stat
 
 
+def check_support_multiagent(alg, config):
+    register_env("multi_mountaincar", lambda _: MultiMountainCar(2))
+    register_env("multi_cartpole", lambda _: MultiCartpole(2))
+    if alg == "DDPG":
+        a = get_agent_class(alg)(config=config, env="multi_mountaincar")
+    else:
+        a = get_agent_class(alg)(config=config, env="multi_cartpole")
+    try:
+        a.train()
+    finally:
+        a.stop()
+
+
 class ModelSupportedSpaces(unittest.TestCase):
+    def setUp(self):
+        ray.init(num_cpus=4)
+
+    def tearDown(self):
+        ray.shutdown()
+
     def testAll(self):
-        ray.init()
         stats = {}
         check_support("IMPALA", {"num_gpus": 0}, stats)
         check_support("DDPG", {"timesteps_per_iteration": 1}, stats)
@@ -137,6 +156,27 @@ class ModelSupportedSpaces(unittest.TestCase):
                   stat)
         self.assertEqual(num_unexpected_errors, 0)
 
+    def testMultiAgent(self):
+        check_support_multiagent("IMPALA", {"num_gpus": 0})
+        check_support_multiagent("DQN", {"timesteps_per_iteration": 1})
+        check_support_multiagent("A3C", {
+            "num_workers": 1,
+            "optimizer": {
+                "grads_per_step": 1
+            }
+        })
+        check_support_multiagent(
+            "PPO", {
+                "num_workers": 1,
+                "num_sgd_iter": 1,
+                "train_batch_size": 10,
+                "sample_batch_size": 10,
+                "sgd_minibatch_size": 1,
+                "simple_optimizer": True,
+            })
+        check_support_multiagent("PG", {"num_workers": 1, "optimizer": {}})
+        check_support_multiagent("DDPG", {"timesteps_per_iteration": 1})
+
 
 if __name__ == "__main__":
     if len(sys.argv) > 1 and sys.argv[1] == "--smoke":