Repository: ray-project/ray

Issue:
[rllib] TorchDiagGaussian doesn’t handle multiple actions correctly.
This is not a contribution.

Ray version: 0.8.2
Python version: 3.6.8
Pytorch version: 1.4
OS: Ubuntu 18.04 Docker

TorchDiagGaussian doesn’t handle multiple actions correctly. As a result, training PPO with Pytorch will crash when the action space has more than 1 action. Here’s minimal reproduction script:

```python
import gym
from gym.spaces import Box
from ray import tune

class ContinuousEnv(gym.Env):
   def __init__(self, config):
       self.action_space = Box(0.0, 1.0, shape=(2,))
       self.observation_space = Box(0.0, 1.0, shape=(1, ))

   def reset(self):
       return [0.0]

   def step(self, action):
       return [0.0], 1.0, False, {}

tune.run(
   "PPO",
   config={"env": ContinuousEnv, "use_pytorch": True, "num_workers": 1})
```



Discussion:
Thanks for filing this. Taking a look ...

Resolving patch:
diff --git a/rllib/agents/ppo/ppo_torch_policy.py b/rllib/agents/ppo/ppo_torch_policy.py
--- a/rllib/agents/ppo/ppo_torch_policy.py
+++ b/rllib/agents/ppo/ppo_torch_policy.py
@@ -67,9 +67,15 @@ def __init__(self,
             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,
diff --git a/rllib/models/torch/torch_action_dist.py b/rllib/models/torch/torch_action_dist.py
--- a/rllib/models/torch/torch_action_dist.py
+++ b/rllib/models/torch/torch_action_dist.py
@@ -71,7 +71,15 @@ def deterministic_sample(self):
 
     @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)


Test patch:
diff --git a/rllib/tests/test_supported_spaces.py b/rllib/tests/test_supported_spaces.py
--- a/rllib/tests/test_supported_spaces.py
+++ b/rllib/tests/test_supported_spaces.py
@@ -11,6 +11,8 @@
 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 @@ def tearDown(self):
         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 @@ def test_impala(self):
         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)