Source code for nnabla_rl.distributions.distribution

# Copyright 2020,2021 Sony Corporation.
# Copyright 2021,2022,2023 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from abc import ABCMeta, abstractmethod
from typing import Optional, Tuple, Union

import numpy as np

import nnabla as nn


[docs]class Distribution(metaclass=ABCMeta):
[docs] @abstractmethod def sample(self, noise_clip: Optional[Tuple[float, float]] = None) -> Union[nn.Variable, np.ndarray]: """Sample a value from the distribution. If noise_clip is specified, the sampled value will be clipped in the given range. Applicability of noise_clip depends on underlying implementation. Args: noise_clip(Tuple[float, float], optional): float tuple of size 2 which contains the min and max value of the noise. Returns: Union[nn.Variable, np.ndarray]: Sampled value """ raise NotImplementedError
@property def ndim(self) -> int: """The number of dimensions of the distribution.""" raise NotImplementedError
[docs] def sample_multiple(self, num_samples: int, noise_clip: Optional[Tuple[float, float]] = None ) -> Union[nn.Variable, np.ndarray]: """Sample mutiple value from the distribution New axis will be added between the first and second axis. Thefore, the returned value shape for mean and variance with shape (batch_size, data_shape) will be changed to (batch_size, num_samples, data_shape) If noise_clip is specified, sampled values will be clipped in the given range. Applicability of noise_clip depends on underlying implementation. Args: num_samples(int): number of samples per batch noise_clip(Tuple[float, float], optional): float tuple of size 2 which contains the min and max value of the noise. Returns: Union[nn.Variable, np.ndarray]: Sampled value. """ raise NotImplementedError
[docs] def choose_probable(self) -> Union[nn.Variable, np.ndarray]: """Compute the most probable action of the distribution. Returns: Union[nn.Variable, np.ndarray]: Probable action of the distribution """ raise NotImplementedError
[docs] def mean(self) -> Union[nn.Variable, np.ndarray]: """Compute the mean of the distribution (if exist) Returns: Union[nn.Variable, np.ndarray]: mean of the distribution Raises: NotImplementedError: The distribution does not have mean """ raise NotImplementedError
[docs] def log_prob(self, x: Union[nn.Variable, np.ndarray]) -> Union[nn.Variable, np.ndarray]: """Compute the log probability of given input. Args: x (Union[nn.Variable, np.ndarray]): Target value to compute the log probability Returns: Union[nn.Variable, np.ndarray]: Log probability of given input """ raise NotImplementedError
[docs] def sample_and_compute_log_prob(self, noise_clip: Optional[Tuple[float, float]] = None) \ -> Union[Tuple[nn.Variable, nn.Variable], Tuple[np.ndarray, np.ndarray]]: """Sample a value from the distribution and compute its log probability. Args: noise_clip(Tuple[float, float], optional): float tuple of size 2 which contains the min and max value of the noise. Returns: Union[Tuple[nn.Variable, nn.Variable], Tuple[np.ndarray, np.ndarray]]: Sampled value and its log probabilty """ raise NotImplementedError
[docs] def entropy(self) -> Union[nn.Variable, np.ndarray]: """Compute the entropy of the distribution. Returns: Union[nn.Variable, np.ndarray]: Entropy of the distribution """ raise NotImplementedError
[docs] def kl_divergence(self, q: 'Distribution') -> Union[nn.Variable, np.ndarray]: """Compute the kullback leibler divergence between given distribution. This function will compute KL(self||q) Args: q(nnabla_rl.distributions.Distribution): target distribution to compute the kl_divergence Returns: Union[nn.Variable, np.ndarray]: Kullback leibler divergence Raises: ValueError: target distribution's type does not match with current distribution type. """ raise NotImplementedError
class DiscreteDistribution(Distribution): pass class ContinuosDistribution(Distribution): pass