From f08687f5503b1a59b1f5ce9a4cb0d6983d4d1554 Mon Sep 17 00:00:00 2001 From: Sven Mika Date: Sun, 8 Mar 2020 21:03:18 +0100 Subject: [PATCH] [RLlib] `rllib train` crashes when using torch PPO/PG/A2C. (#7508) * Fix. * Rollback. * TEST. * TEST. * TEST. * TEST. * TEST. * TEST. * TEST. * TEST. * TEST. * TEST. * TEST. * TEST. * TEST. * TEST. * TEST. * TEST. * TEST. * TEST. * TEST. * TEST. * TEST. * TEST. * TEST. --- rllib/models/torch/torch_action_dist.py | 6 ++++-- rllib/policy/torch_policy.py | 8 +++++--- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/rllib/models/torch/torch_action_dist.py b/rllib/models/torch/torch_action_dist.py index f3d6f03e7..d8fd9e7f6 100644 --- a/rllib/models/torch/torch_action_dist.py +++ b/rllib/models/torch/torch_action_dist.py @@ -12,7 +12,8 @@ class TorchDistributionWrapper(ActionDistribution): @override(ActionDistribution) def __init__(self, inputs, model): - inputs = torch.Tensor(inputs) + if not isinstance(inputs, torch.Tensor): + inputs = torch.Tensor(inputs) super().__init__(inputs, model) # Store the last sample here. self.last_sample = None @@ -46,7 +47,8 @@ class TorchCategorical(TorchDistributionWrapper): @override(ActionDistribution) def __init__(self, inputs, model=None, temperature=1.0): assert temperature > 0.0, "Categorical `temperature` must be > 0.0!" - super().__init__(inputs / temperature, model) + inputs /= temperature + super().__init__(inputs, model) self.dist = torch.distributions.categorical.Categorical( logits=self.inputs) diff --git a/rllib/policy/torch_policy.py b/rllib/policy/torch_policy.py index 4972ef898..417107fbe 100644 --- a/rllib/policy/torch_policy.py +++ b/rllib/policy/torch_policy.py @@ -7,6 +7,7 @@ from ray.rllib.policy.sample_batch import SampleBatch from ray.rllib.utils.annotations import override, DeveloperAPI from ray.rllib.utils.framework import try_import_torch from ray.rllib.utils.schedules import ConstantSchedule, PiecewiseSchedule +from ray.rllib.utils.torch_ops import convert_to_non_torch_type from ray.rllib.utils.tracking_dict import UsageTrackingDict torch, _ = try_import_torch() @@ -94,12 +95,13 @@ class TorchPolicy(Policy): extra_action_out = self.extra_action_out(input_dict, state_batches, self.model, action_dist) if logp is not None: + logp = convert_to_non_torch_type(logp) extra_action_out.update({ - ACTION_PROB: torch.exp(logp), + ACTION_PROB: np.exp(logp), ACTION_LOGP: logp }) - return (actions.cpu().numpy(), [h.cpu().numpy() for h in state], - extra_action_out) + return convert_to_non_torch_type( + (actions, state, extra_action_out)) @override(Policy) def compute_log_likelihoods(self,