diff --git a/rllib/models/tf/tf_action_dist.py b/rllib/models/tf/tf_action_dist.py index 27813ddd3..06c182372 100644 --- a/rllib/models/tf/tf_action_dist.py +++ b/rllib/models/tf/tf_action_dist.py @@ -335,6 +335,14 @@ class SquashedGaussian(TFActionDistribution): axis=-1) return log_prob + @override(ActionDistribution) + def entropy(self) -> TensorType: + raise ValueError("Entropy not defined for SquashedGaussian!") + + @override(ActionDistribution) + def kl(self, other: ActionDistribution) -> TensorType: + raise ValueError("KL not defined for SquashedGaussian!") + def _squash(self, raw_values: TensorType) -> TensorType: # Returned values are within [low, high] (including `low` and `high`). squashed = ((tf.math.tanh(raw_values) + 1.0) / 2.0) * \ diff --git a/rllib/models/torch/torch_action_dist.py b/rllib/models/torch/torch_action_dist.py index 83f900a5e..ecc8aa276 100644 --- a/rllib/models/torch/torch_action_dist.py +++ b/rllib/models/torch/torch_action_dist.py @@ -243,6 +243,14 @@ class TorchSquashedGaussian(TorchDistributionWrapper): torch.log(1 - unsquashed_values_tanhd**2 + SMALL_NUMBER), dim=-1) return log_prob + @override(TorchDistributionWrapper) + def entropy(self) -> TensorType: + raise ValueError("Entropy not defined for SquashedGaussian!") + + @override(TorchDistributionWrapper) + def kl(self, other: ActionDistribution) -> TensorType: + raise ValueError("KL not defined for SquashedGaussian!") + def _squash(self, raw_values: TensorType) -> TensorType: # Returned values are within [low, high] (including `low` and `high`). squashed = ((torch.tanh(raw_values) + 1.0) / 2.0) * \