From d8eeb9641314740572e81f9836cbce3e5b8f2b73 Mon Sep 17 00:00:00 2001 From: Sven Mika Date: Mon, 2 Mar 2020 19:53:19 +0100 Subject: [PATCH] Fix issue with torch PPO not handling action spaces of shape=(>1,). (#7398) --- rllib/agents/ppo/ppo_torch_policy.py | 14 +++-- rllib/models/torch/torch_action_dist.py | 10 +++- rllib/tests/test_supported_spaces.py | 77 +++++++++++++++---------- 3 files changed, 64 insertions(+), 37 deletions(-) diff --git a/rllib/agents/ppo/ppo_torch_policy.py b/rllib/agents/ppo/ppo_torch_policy.py index af6cd0ec5..820963c93 100644 --- a/rllib/agents/ppo/ppo_torch_policy.py +++ b/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, diff --git a/rllib/models/torch/torch_action_dist.py b/rllib/models/torch/torch_action_dist.py index a1f7748e9..fd6717240 100644 --- a/rllib/models/torch/torch_action_dist.py +++ b/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) diff --git a/rllib/tests/test_supported_spaces.py b/rllib/tests/test_supported_spaces.py index a09d2e5b3..6a732042b 100644 --- a/rllib/tests/test_supported_spaces.py +++ b/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)