Fix issue with torch PPO not handling action spaces of shape=(>1,). (#7398)

rllib/agents/ppo/ppo_torch_policy.py

@@ -67,9 +67,15 @@ class PPOLoss:
             vf_loss_coeff (float): Coefficient of the value function loss
             use_gae (bool): If true, use the Generalized Advantage Estimator.
         """
+        if valid_mask is not None:
 
-        def reduce_mean_valid(t):
-            return torch.mean(t * valid_mask)
+            def reduce_mean_valid(t):
+                return torch.mean(t * valid_mask)
+
+        else:
+
+            def reduce_mean_valid(t):
+                return torch.mean(t)
 
         prev_dist = dist_class(prev_logits, model)
         # Make loss functions.
@@ -109,13 +115,11 @@ def ppo_surrogate_loss(policy, model, dist_class, train_batch):
     logits, state = model.from_batch(train_batch)
     action_dist = dist_class(logits, model)
 
+    mask = None
     if state:
         max_seq_len = torch.max(train_batch["seq_lens"])
         mask = sequence_mask(train_batch["seq_lens"], max_seq_len)
         mask = torch.reshape(mask, [-1])
-    else:
-        mask = torch.ones_like(
-            train_batch[Postprocessing.ADVANTAGES], dtype=torch.bool)
 
     policy.loss_obj = PPOLoss(
         dist_class,

rllib/models/torch/torch_action_dist.py

@@ -71,7 +71,15 @@ class TorchDiagGaussian(TorchDistributionWrapper):
 
     @override(TorchDistributionWrapper)
     def logp(self, actions):
-        return TorchDistributionWrapper.logp(self, actions).sum(-1)
+        return super().logp(actions).sum(-1)
+
+    @override(TorchDistributionWrapper)
+    def entropy(self):
+        return super().entropy().sum(-1)
+
+    @override(TorchDistributionWrapper)
+    def kl(self, other):
+        return super().kl(other).sum(-1)
 
     @staticmethod
     @override(ActionDistribution)

rllib/tests/test_supported_spaces.py

@@ -11,6 +11,8 @@ from ray.rllib.utils.framework import try_import_tf
 from ray.rllib.agents.registry import get_agent_class
 from ray.rllib.models.tf.fcnet_v2 import FullyConnectedNetwork as FCNetV2
 from ray.rllib.models.tf.visionnet_v2 import VisionNetwork as VisionNetV2
+from ray.rllib.models.torch.visionnet import VisionNetwork as TorchVisionNetV2
+from ray.rllib.models.torch.fcnet import FullyConnectedNetwork as TorchFCNetV2
 from ray.rllib.tests.test_multi_agent_env import MultiCartpole, \
     MultiMountainCar
 from ray.rllib.utils.error import UnsupportedSpaceException
@@ -75,10 +77,11 @@ def check_support(alg, config, stats, check_bounds=False, name=None):
     covered_o = set()
     config["log_level"] = "ERROR"
     first_error = None
+    torch = config.get("use_pytorch", False)
     for a_name, action_space in ACTION_SPACES_TO_TEST.items():
         for o_name, obs_space in OBSERVATION_SPACES_TO_TEST.items():
-            print("=== Testing {} A={} S={} ===".format(
-                alg, action_space, obs_space))
+            print("=== Testing {} (torch={}) A={} S={} ===".format(
+                alg, torch, action_space, obs_space))
             stub_env = make_stub_env(action_space, obs_space, check_bounds)
             register_env("stub_env", lambda c: stub_env())
             stat = "ok"
@@ -86,14 +89,26 @@ def check_support(alg, config, stats, check_bounds=False, name=None):
             try:
                 if a_name in covered_a and o_name in covered_o:
                     stat = "skip"  # speed up tests by avoiding full grid
+                # TODO(sven): Add necessary torch distributions.
+                elif torch and a_name in ["tuple", "multidiscrete"]:
+                    stat = "unsupported"
                 else:
                     a = get_agent_class(alg)(config=config, env="stub_env")
                     if alg not in ["DDPG", "ES", "ARS", "SAC"]:
                         if o_name in ["atari", "image"]:
-                            assert isinstance(a.get_policy().model,
-                                              VisionNetV2)
+                            if torch:
+                                assert isinstance(
+                                    a.get_policy().model, TorchVisionNetV2)
+                            else:
+                                assert isinstance(
+                                    a.get_policy().model, VisionNetV2)
                         elif o_name in ["vector", "vector2"]:
-                            assert isinstance(a.get_policy().model, FCNetV2)
+                            if torch:
+                                assert isinstance(
+                                    a.get_policy().model, TorchFCNetV2)
+                            else:
+                                assert isinstance(
+                                    a.get_policy().model, FCNetV2)
                     a.train()
                     covered_a.add(a_name)
                     covered_o.add(o_name)
@@ -144,15 +159,15 @@ class ModelSupportedSpaces(unittest.TestCase):
         ray.shutdown()
 
     def test_a3c(self):
-        check_support(
-            "A3C", {
-                "num_workers": 1,
-                "optimizer": {
-                    "grads_per_step": 1
-                }
-            },
-            self.stats,
-            check_bounds=True)
+        config = {
+            "num_workers": 1,
+            "optimizer": {
+                "grads_per_step": 1
+            }
+        }
+        check_support("A3C", config, self.stats, check_bounds=True)
+        config["use_pytorch"] = True
+        check_support("A3C", config, self.stats, check_bounds=True)
 
     def test_appo(self):
         check_support("APPO", {"num_gpus": 0, "vtrace": False}, self.stats)
@@ -201,25 +216,25 @@ class ModelSupportedSpaces(unittest.TestCase):
         check_support("IMPALA", {"num_gpus": 0}, self.stats)
 
     def test_ppo(self):
-        check_support(
-            "PPO", {
-                "num_workers": 1,
-                "num_sgd_iter": 1,
-                "train_batch_size": 10,
-                "sample_batch_size": 10,
-                "sgd_minibatch_size": 1,
-            },
-            self.stats,
-            check_bounds=True)
+        config = {
+            "num_workers": 1,
+            "num_sgd_iter": 1,
+            "train_batch_size": 10,
+            "sample_batch_size": 10,
+            "sgd_minibatch_size": 1,
+        }
+        check_support("PPO", config, self.stats, check_bounds=True)
+        config["use_pytorch"] = True
+        check_support("PPO", config, self.stats, check_bounds=True)
 
     def test_pg(self):
-        check_support(
-            "PG", {
-                "num_workers": 1,
-                "optimizer": {}
-            },
-            self.stats,
-            check_bounds=True)
+        config = {
+            "num_workers": 1,
+            "optimizer": {}
+        }
+        check_support("PG", config, self.stats, check_bounds=True)
+        config["use_pytorch"] = True
+        check_support("PG", config, self.stats, check_bounds=True)
 
     def test_sac(self):
         check_support("SAC", {}, self.stats, check_bounds=True)