From 2b1b255a67ff38a10f22d18446d020197d84a53f Mon Sep 17 00:00:00 2001 From: Brian Delhaisse Date: Sat, 30 Mar 2019 23:53:46 +0100 Subject: [PATCH] update distribution modules/layers --- .gitignore | 1 - pyrobolearn/backends/README.md | 15 +- pyrobolearn/backends/torch_backend.py | 118 +++- pyrobolearn/distributions/README.md | 17 +- pyrobolearn/distributions/__init__.py | 8 + pyrobolearn/distributions/bernoulli.py | 46 ++ pyrobolearn/distributions/categorical.py | 47 ++ pyrobolearn/distributions/gaussian.py | 4 + pyrobolearn/distributions/modules.py | 731 +++++++++++++++++++++++ pyrobolearn/models/gaussian.py | 52 +- 10 files changed, 1006 insertions(+), 33 deletions(-) create mode 100644 pyrobolearn/distributions/bernoulli.py create mode 100644 pyrobolearn/distributions/categorical.py create mode 100644 pyrobolearn/distributions/modules.py diff --git a/.gitignore b/.gitignore index c7e5a28..63c4b3f 100644 --- a/.gitignore +++ b/.gitignore @@ -109,4 +109,3 @@ venv.bak/ .ipynb_checkpoints/ .idea/ 0_VRDemoSettings.txt -tests/ diff --git a/pyrobolearn/backends/README.md b/pyrobolearn/backends/README.md index 5e95153..a5528b6 100644 --- a/pyrobolearn/backends/README.md +++ b/pyrobolearn/backends/README.md @@ -1,8 +1,17 @@ ## Backends -The general idea is to provide different backends such that different tensor frameworks can be used. These include for instance numpy (with autograd), pytorch, and tensorflow. -These frameworks use different data structures and different signatures for the methods. +The general idea is to provide different backends such that different tensor frameworks can be used. These include for +instance numpy (with autograd), pytorch, and tensorflow. These frameworks use different data structures and different +method signatures. By defining a common API, it would ease the use of these various frameworks as the syntax +would be the same. For instance, in numpy, the outer product between two arrays is performed using `np.outer` +while in pytorch it is carried out by calling `torch.ger`. Defining a common API would solve these issues. +Also, using backends, we could convert inside each function the given data to the appropriate data structure. +For instance, a pytorch function defined in the backend could easily accept a `np.array` as input and convert it +automatically to a `torch.Tensor`. It would be nice to have learning models that are more or less independent of the tensor framework as done in Keras. -This is mostly an idea that I had in a later stage, and thus is not operational for the moment. It would require to refactor a bit the code, as currently our code is coupled to the pytorch and numpy frameworks. +This is mostly an idea that I had in a later stage, and thus is not operational for the moment. It would require to +refactor a bit the code, as currently our code is coupled to the pytorch and numpy frameworks. Also, it would +require to provide the same functionalities in the various frameworks, and thus implement their missing +functionalities. diff --git a/pyrobolearn/backends/torch_backend.py b/pyrobolearn/backends/torch_backend.py index 3322b58..8d43fa7 100644 --- a/pyrobolearn/backends/torch_backend.py +++ b/pyrobolearn/backends/torch_backend.py @@ -1,6 +1,80 @@ +#!/usr/bin/env python +"""Provide the torch backend API. +""" import torch from torch import * +import numpy as np +import tensorflow as tf + + +__author__ = "Brian Delhaisse" +__copyright__ = "Copyright 2018, PyRoboLearn" +__credits__ = ["Brian Delhaisse"] +__license__ = "MIT" +__version__ = "1.0.0" +__maintainer__ = "Brian Delhaisse" +__email__ = "briandelhaisse@gmail.com" +__status__ = "Development" + + +# define decorator that converts the given data structure (numpy array, tensorflow tensor, or torch tensor) to a torch +# tensor. +def to_torch(function): + """ + Decorator around a given function. + + Args: + function (callable): function to decorate / wrap. + + Returns: + callable: function that wraps the initial function by making sure its argument is a torch tensor. + """ + + def wrapper(data, *args, **kwargs): + """Process the given argument. + + Args: + data (np.array, torch.Tensor, tf.Tensor): input data. + """ + if not isinstance(data, (tuple, list)): + data = [data] + + d = [] + for datum in data: + # convert to torch Tensor + if isinstance(datum, torch.Tensor): # if torch tensor, do nothing + pass + elif isinstance(datum, np.ndarray): # if numpy array, convert to torch tensor + datum = torch.from_numpy(datum).float() + elif isinstance(datum, tf.Tensor): # if tensorflow tensor, convert to torch tensor + # run the tensorflow session + sess = tf.Session() + with sess.as_default(): + datum = datum.eval() + + # convert the numpy array to torch tensor + if not isinstance(datum, np.ndarray): + raise TypeError("The data returned by evaluating the tensorflow session, is not a numpy array, " + "but instead: {}".format(type(datum))) + datum = torch.from_numpy(datum).float() + else: + raise TypeError("Expecting the input to be a `torch.Tensor`, `np.ndarray`, or `tf.Tensor`, instead " + "got: {}".format(type(datum))) + d.append(datum) + + if len(d) == 1: + data = d[0] + else: + data = d + + # call inner function on the given argument + data = function(data, *args, **kwargs) + + # return torch Tensor + return data + + return wrapper def array(data, dtype=None, copy=True, device=None, requires_grad=False, ndmin=0): @@ -11,8 +85,50 @@ def inv(data, out=None): return torch.inverse(data, out=out) +@to_torch def concatenate(data, axis=0, out=None): - return torch.cat(data, axis, out) + return torch.cat(data, dim=axis, out=out) # np.size vs torch.nelement() vs torch.size() +# TODO +# torch.ger --> torch.outer +# torch.dot only works on 1D vector compared to np.dot, in torch, need to use torch.mm +# missing torch.dstack +# missing torch.vstack +# missing torch.hstack + + +# Tests +if __name__ == '__main__': + + # define function to evaluate tf tensor + def tf_eval(tensor): + sess = tf.Session() + with sess.as_default(): + tensor = tensor.eval() + return tensor + + # create 3 variables + a = np.ones(2) + b = torch.ones(2) + c = tf.ones(2) + + # concatenate + numpy_result = np.concatenate((a, a)) + torch_result = torch.cat((b, b)) + tf_result = tf_eval(tf.concat((c, c))) + + print("np.concatenate: {}".format(numpy_result)) + print("torch.cat: {}".format(torch_result)) + print("tf.concat: {}".format(tf_result)) + + result_1 = concatenate((a, a)) + result_2 = concatenate((b, b)) + result_3 = concatenate((c, c)) + result_4 = concatenate((a, b, c)) + + print("concatenate two np.array: {}".format(result_1)) + print("concatenate two torch.Tensor: {}".format(result_2)) + print("concatenate two tf.Tensor: {}".format(result_3)) + print("concatenate one np.array, one torch.Tensor, one tf.Tensor".format(result_4)) diff --git a/pyrobolearn/distributions/README.md b/pyrobolearn/distributions/README.md index 9e2d614..1197d4e 100644 --- a/pyrobolearn/distributions/README.md +++ b/pyrobolearn/distributions/README.md @@ -1,4 +1,17 @@ -## Probability distributions +## Probability distributions and layers -This folder mainly contains wrappers to the various `torch.distributions.*`, and sometimes extend few of them. +This folder mainly contains wrappers to the various `torch.distributions.*` and provides few additional features / +functionalities. It also provides several `torch.nn.Module` layers/modules that accepts as input the base output +or the output of a learning model (such as a linear model, or multilayer perceptron) and returns a distribution +defined on the output of such layers/modules. +For instance, you can defined a Gaussian distribution module/layer as such: +```python +from pyrobolearn.distributions.modules import * + +mean = MeanModule(num_inputs=10, num_outputs=5) +covariance = FullCovarianceModule(num_inputs=10, num_outputs=5) +gaussian = GaussianModule(mean=mean, covariance=covariance) +probs = gaussian(base_output) # this will feed the `base_output` to the previously defined mean and covariance, + # and will returned a Gaussian distribution based on their outputs. +``` diff --git a/pyrobolearn/distributions/__init__.py b/pyrobolearn/distributions/__init__.py index e69de29..f03ed78 100644 --- a/pyrobolearn/distributions/__init__.py +++ b/pyrobolearn/distributions/__init__.py @@ -0,0 +1,8 @@ + +# import distributions +from .bernoulli import Bernoulli +from .categorical import Categorical +from .gaussian import Gaussian + +# import modules / layers +from .modules import * diff --git a/pyrobolearn/distributions/bernoulli.py b/pyrobolearn/distributions/bernoulli.py new file mode 100644 index 0000000..ef5e4e9 --- /dev/null +++ b/pyrobolearn/distributions/bernoulli.py @@ -0,0 +1,46 @@ +#!/usr/bin/env python +"""Define the discrete Bernoulli distribution class. +""" + +import torch +import numpy as np + + +__author__ = "Brian Delhaisse" +__copyright__ = "Copyright 2018, PyRoboLearn" +__credits__ = ["Brian Delhaisse"] +__license__ = "MIT" +__version__ = "1.0.0" +__maintainer__ = "Brian Delhaisse" +__email__ = "briandelhaisse@gmail.com" +__status__ = "Development" + + +class Bernoulli(torch.distributions.Bernoulli): + r"""Bernoulli distribution + + Type: discrete, binary + + "The Bernoulli distribution is the discrete probability distribution of a random variable which takes the value 1 + with probability :math:`p` and the value 0 with probability :math:`q = 1-p`, that is, the probability distribution + of any single experiment that asks a yes/no question; the question results in a boolean-valued outcome, a single + bit of information whose value is success with probability :math:`p` and failure with probability :math:`q`." [1] + + References: + [1] Bernoulli distribution: https://en.wikipedia.org/wiki/Bernoulli_distribution + """ + + def __init__(self, probs=None, logits=None): + """ + Initialize the Bernoulli distribution on the given manifold. + + Args: + probs (torch.Tensor, None): event probabilities module. + logits (torch.Tensor, None): event logits module. + """ + # call superclass + super(Bernoulli, self).__init__(probs=probs, logits=logits) + + def mode(self): + """Return the mode of the Bernoulli distribution.""" + return torch.gt(self.probs, 0.5).float() diff --git a/pyrobolearn/distributions/categorical.py b/pyrobolearn/distributions/categorical.py new file mode 100644 index 0000000..d767db3 --- /dev/null +++ b/pyrobolearn/distributions/categorical.py @@ -0,0 +1,47 @@ +#!/usr/bin/env python +"""Define the discrete Categorical distribution class. +""" + +import torch + +__author__ = "Brian Delhaisse" +__copyright__ = "Copyright 2018, PyRoboLearn" +__credits__ = ["Brian Delhaisse"] +__license__ = "MIT" +__version__ = "1.0.0" +__maintainer__ = "Brian Delhaisse" +__email__ = "briandelhaisse@gmail.com" +__status__ = "Development" + + +class Categorical(torch.distributions.Categorical): + r"""Categorical distribution + + Type: discrete, multiple categories + + The Categorical module accepts as inputs the discrete logits or probabilities modules, and returns the categorical + distribution (that inherits from `torch.distributions.Categorical`). + + Description: "A categorical distribution (also called a generalized Bernoulli distribution, multinoulli + distribution) is a discrete probability distribution that describes the possible results of a random variable that + can take on one of K possible categories, with the probability of each category separately specified." [1] + + References: + [1] Categorical distribution: https://en.wikipedia.org/wiki/Categorical_distribution + """ + + def __init__(self, probs=None, logits=None): + """ + Initialize the Categorical distribution on the given manifold. + + Args: + probs (torch.Tensor, None): event probabilities module. + logits (torch.Tensor, None): event logits module. + """ + # call superclass + super(Categorical, self).__init__(probs=probs, logits=logits) + + @property + def mode(self): + """Return the mode of the Categorical distribution.""" + return self.probs.argmax(dim=-1, keepdim=True) diff --git a/pyrobolearn/distributions/gaussian.py b/pyrobolearn/distributions/gaussian.py index 6ee391d..9caa746 100644 --- a/pyrobolearn/distributions/gaussian.py +++ b/pyrobolearn/distributions/gaussian.py @@ -4,6 +4,10 @@ This distribution is so important in the field of Machine Learning that we extended the functionalities of the basic `torch.distributions.MultivariateNormal`. This distribution will notably be used for Gaussian Mixture Models, Probabilistic Movement Primitives, Kernelized Movement Primitives, etc. + +References: + [1] torch.distributions: https://pytorch.org/docs/stable/distributions.html + [2] Gaussian distribution: pyrobolearn/models/gaussian """ import torch diff --git a/pyrobolearn/distributions/modules.py b/pyrobolearn/distributions/modules.py new file mode 100644 index 0000000..d6c4ada --- /dev/null +++ b/pyrobolearn/distributions/modules.py @@ -0,0 +1,731 @@ +#!/usr/bin/env python +"""Provide the various common probability distribution layers / modules. + +This file provides layers / modules that can output probability distributions that inherit from +`torch.distributions.*`. Several pieces of code were inspired from [1, 2]. + +References: + [1] torch.distributions: https://pytorch.org/docs/stable/distributions.html + [2] pytorch-a2c-ppo-acktr: + - https://github.com/ikostrikov/pytorch-a2c-ppo-acktr-gail/blob/master/a2c_ppo_acktr/utils.py + - https://github.com/ikostrikov/pytorch-a2c-ppo-acktr/blob/master/distributions.py + [3] Gaussian distribution: pyrobolearn/models/gaussian +""" + +from abc import ABCMeta + +import numpy as np +import torch +from gaussian import Gaussian as GaussianDistribution +from categorical import Categorical as CategoricalDistribution +from bernoulli import Bernoulli as BernoulliDistribution + + +__author__ = "Brian Delhaisse" +__copyright__ = "Copyright 2018, PyRoboLearn" +__credits__ = ["Brian Delhaisse", "Ilya Kostrikov"] +__license__ = "MIT" +__version__ = "1.0.0" +__maintainer__ = "Brian Delhaisse" +__email__ = "briandelhaisse@gmail.com" +__status__ = "Development" + + +def init_module(module, init_weight=None, init_bias=None): + """ + Initialize the given module using the given initialization scheme for the weight and bias terms. + + The user can select the initialization scheme from `torch.nn.init`. This function is taken from [1] and modified. + + Args: + module (torch.nn.Module): torch module to initialize. + init_weight (callable, None): this is a callable function that accepts as input the module to initialize its + weights. If None, it will leave the module weights untouched. + init_bias (callable, None): this is a callable function that accepts as input the module to initialize its + bias terms. If None, it will leave the module bias untouched. + + Returns: + torch.nn.Module: initialized module. + + References: + [1] https://github.com/ikostrikov/pytorch-a2c-ppo-acktr-gail/blob/master/a2c_ppo_acktr/utils.py + """ + if init_weight is not None: + init_weight(module.weight.data) + if init_bias is not None: + init_bias(module.bias.data) + return module + + +def wrap_init_tensor(init_tensor, *args, **kwargs): + r"""Define a higher order function that accepts as inputs the initial method to initialize a tensor, as well + as its arguments, and returns a function that only awaits for its tensor input. With this, you can use the above + :func:`init_module` function quite easily. For instance: + + Examples: + >>> module = torch.nn.Linear(in_features=10, out_features=5) + >>> weight_init = wrap_init_tensor(torch.nn.init.orthogonal_, gain=1.) + >>> weight_bias = wrap_init_tensor(torch.nn.init.constant_, val=0) + >>> module = init_module(module, wrap_init_tensor(module, weight_init, weight_bias)) + + Returns: + callable: return the callable function that only accepts a tensor as input. + """ + def init(tensor): + return init_tensor(tensor, *args, **kwargs) + return init + + +def init_orthogonal_weights_and_constant_biases(module, gain=1., val=0.): + """Initialize the weights of the module to be orthogonal and with a bias of 0s. This is inspired by [1]. + + Args: + gain (float): optional scaling factor for the orthogonal weights. + val (float): val: the value to fill the bias tensor with. + + Returns: + torch.nn.Module: the initialized module. + + References: + [1] https://github.com/ikostrikov/pytorch-a2c-ppo-acktr-gail/blob/master/a2c_ppo_acktr/distributions.py + """ + weight_init = wrap_init_tensor(torch.nn.init.orthogonal_, gain=gain) + weight_bias = wrap_init_tensor(torch.nn.init.constant_, val=val) + module = init_module(module, wrap_init_tensor(module, weight_init, weight_bias)) + return module + + +class DistributionModule(torch.nn.Module): + r"""Probability distribution module + + Define a wrapper around `torch.distributions.Distribution` with lazy creation and additional features. + Everytime this object is called given an input it wraps that given input, and return a distribution over it. + """ + + def __init__(self, distribution): + """ + Initialize the distribution. + + Args: + distribution (type): subclass of `torch.distributions.Distribution`. + """ + super(DistributionModule, self).__init__() + self.distribution = distribution + + @property + def distribution(self): + """Return the distribution type.""" + return self._distribution + + @distribution.setter + def distribution(self, distribution): + """Set the distribution.""" + if not issubclass(distribution, torch.distributions.Distribution): + raise TypeError("Expecting the given distribution to be a subclass of `torch.distributions.Distribution`, " + "instead got: {}".format(distribution)) + self._distribution = distribution + + def forward(self, x): + """Forward the given inputs :attr:`x`.""" + raise NotImplementedError + + +class FixedVectorModule(torch.nn.Module): + r"""Fixed vector generator (module) + + Generate the fixed vector. This generates a fixed vector everytime it is called. + If N samples are given at the input such that it has a shape (N,I), it returns N copy of the vector of shape (N,M). + """ + + def __init__(self, vector): + """ + Initialize the fixed generator. + + Args: + vector (torch.Tensor): fixed vector + """ + super(FixedVectorModule, self).__init__() + if not isinstance(vector, torch.Tensor): + raise TypeError("Expecting the vector to be an instance of `torch.Tensor`, instead got: " + "{}".format(type(vector))) + self._vector = vector + + def forward(self, x): + """Take as input the base output vector / matrix and return the diagonal covariance matrix(ces). + + Args: + x (torch.Tensor): base output vector / matrix of shape (N, B) + + Returns: + torch.Tensor: vector / matrix of shape (N, D) + """ + return self._vector.repeat(x.size(0), 1) + + +class VectorModule(torch.nn.Module): + r"""Vector generator (module) + + Generate an output vector / matrix given an input vector / matrix. Specifically, it is just a linear module. + """ + + def __init__(self, num_inputs, num_outputs): + """ + Initialize the mean generator. + + Args: + num_inputs (int): size of the base output vector + num_outputs (int): size of the output vector + """ + super(VectorModule, self).__init__() + # linear mapping between the base output vector / matrix and the mean output vector / matrix + model = torch.nn.Linear(in_features=num_inputs, out_features=num_outputs) + self._model = init_orthogonal_weights_and_constant_biases(model) + + def forward(self, x): + """Take as input the base output vector / matrix and return the vector / matrix. + + Args: + x (torch.Tensor): base output vector / matrix of shape (N, B) + + Returns: + torch.Tensor: vector / matrix of shape (N, D) + """ + return self._model(x) + + +class IdentityModule(torch.nn.Module): + r"""Identity Module. + + This just returns whatever tensor is given. + """ + def __init__(self): + """Initialize the identity module.""" + super(IdentityModule, self).__init__() + + def forward(self, x): + """Return the same given input :attr:`x`.""" + return x + + +class FixedMeanModule(torch.nn.Module): # this is the same as FixedVector (but with different documentation) + r"""Fixed mean generator + + Generate the mean of the multivariate Normal distribution. This generates a fixed mean everytime it is called. + If N samples are given at the input such that it has a shape (N,I), it returns N copy of the mean of shape (N,M). + """ + + def __init__(self, mean): + """ + Initialize the mean generator. + + Args: + mean (torch.Tensor): mean vector + """ + super(FixedMeanModule, self).__init__() + if not isinstance(mean, torch.Tensor): + raise TypeError("Expecting the mean vector to be an instance of `torch.Tensor`, instead got: " + "{}".format(type(mean))) + self._mean = mean + + def forward(self, x): + """Take as input the base output vector / matrix and return the mean vector / matrix. + + Args: + x (torch.Tensor): base output vector / matrix of shape (N, B) + + Returns: + torch.Tensor: mean vector / matrix of shape (N, D) + """ + return self._mean.repeat(x.size(0), 1) + + +class MeanModule(torch.nn.Module): # this is the same as VectorModule (but with different documentation) + r"""Mean generator + + Generate the mean of the multivariate Normal distribution. + """ + + def __init__(self, num_inputs, num_outputs): + """ + Initialize the mean generator. + + Args: + num_inputs (int): size of the base output vector + num_outputs (int): size of the action mean vector + """ + super(MeanModule, self).__init__() + # linear mapping between the base output vector / matrix and the mean output vector / matrix + model = torch.nn.Linear(in_features=num_inputs, out_features=num_outputs) + self._model = init_orthogonal_weights_and_constant_biases(model) + + def forward(self, x): + """Take as input the base output vector / matrix and return the mean vector / matrix. + + Args: + x (torch.Tensor): base output vector / matrix of shape (N, B) + + Returns: + torch.Tensor: mean vector / matrix of shape (N, D) + """ + return self._model(x) + + +class FixedDiagonalCovarianceModule(torch.nn.Module): + r"""Fixed Diagonal Covariance Generator + + This covariance receives the mean as input, and generates samples using that mean and a fixed diagonal covariance + matrix set during the instantiation of this class. + + Note that the standard deviations must be non-negatives. + """ + + def __init__(self, variances=None, stddev=None): + """ + Initialize the fixed diagonal covariance matrix generator. + + Args: + variances (torch.Tensor, None): vector of variances. If None, the standard deviations have to be defined. + stddev (torch.Tensor, None): vector of standard deviations. If the variances is None, the standard + deviations will be considered. + """ + super(FixedDiagonalCovarianceModule, self).__init__() + if variances is None: + if stddev is None: + raise ValueError("The variances or the standard deviations have to be specified") + variances = stddev.pow(2) + + # check if negative elements in variances + if torch.any(variances < 0.): + raise ValueError("Expecting the variances to be strictly positive, found some negative variances: " + "{}".format(type(variances))) + + # if variance very close to zero + if torch.any(variances.isclose(torch.tensor(0.))): + # add small offset + variances += 1.e-4 + + # create fixed diagonal covariance matrix + dim = variances.size(-1) + self._covariance = torch.diag(variances).view(1, dim, dim) + + def forward(self, x): + """Take as input the base output vector / matrix and return the fixed diagonal covariance matrix(ces). + + Args: + x (torch.Tensor): base output vector / matrix of shape (N, B) + + Returns: + torch.Tensor: covariance matrices (one covariance matrix for each sample) of shape (N, D, D) + """ + # stack N times the covariance + covariance = self._covariance.repeat(x.size(0), 1, 1) + + # return the fixed diagonal covariance + return covariance + + +class FixedCovarianceModule(torch.nn.Module): + r"""Fixed Covariance Generator + + This Gaussian receives the mean as input, and generates samples using that mean and a fixed covariance matrix set + during the instantiation of this class. + """ + + def __init__(self, covariance=None, precision=None, tril=None): + """ + Initialize the fixed covariance matrix generator. + + Args: + covariance (None, torch.Tensor): positive-definite covariance matrix. + precision (None, torch.Tensor): positive-definite precision matrix. + tril (None, torch.Tensor): lower triangular matrix which is the Cholesky decomposition of the covariance + matrix, with positive-valued diagonal. + """ + super(FixedCovarianceModule, self).__init__() + if covariance is None and precision is None and tril is None: + raise ValueError("The covariance, the precision, or the lower triangular factor of the covariance has to " + "be specified.") + if covariance is not None: + dim = covariance.size(-1) + if precision is not None: + dim = precision.size(-1) + if tril is not None: + dim = tril.size(-1) + + # let the PyTorch framework check if the covariance matrix is correct + distribution = torch.distributions.MultivariateNormal(loc=torch.zeros(dim), covariance_matrix=covariance, + precision_matrix=precision, scale_tril=None) + + # get the covariance matrix + self._covariance = distribution.covariance_matrix + + def forward(self, x): + """Take as input the base output vector / matrix and return the fixed covariance matrix(ces). + + Args: + x (torch.Tensor): base output vector / matrix of shape (N, B) + + Returns: + torch.Tensor: covariance matrices (one covariance matrix for each sample) of shape (N, D, D) + """ + # stack N times the covariance + covariance = self._covariance.repeat(x.size(0), 1, 1) + return covariance + + +class DiagonalCovarianceModule(torch.nn.Module): + r"""Diagonal Covariance Generator + + This class receives generates the diagonal of a covariance matrix given the base output vector / matrix. + Note that a covariance matrix has to be positive semi-definite. In the case of a diagonal matrix, + this is achieved by having the diagonal entries (=the variances) to be non-negative. If the standard deviations + are given they can have any values as we will square them later to form the diagonal entries of the covariance + matrix. By squaring them, they all become non-negative. If instead the variance are given as inputs we have to + make sure that they are non-negative. A quick hack is to give the logarithm of the variance which is always + positive. + """ + + def __init__(self, num_inputs, num_outputs, offset=1.e-4): + """ + Initialize the diagonal covariance generator. + + Args: + num_inputs (int): size of the base output vector + num_outputs (int): size of the action mean vector + offset (float): small offset to be added to the diagonal elements of the covariance matrix such that + it is positive definite. + """ + super(DiagonalCovarianceModule, self).__init__() + # set output dimension + self._dim = int(num_outputs) + + # define diagonal offset such that the covariance is positive definite + self._offset = offset * torch.diag(torch.ones(self._dim)) + + # linear mapping between the base output vector / matrix to the covariance diagonal elements + model = torch.nn.Linear(in_features=num_inputs, out_features=num_outputs) + self._model = init_orthogonal_weights_and_constant_biases(model) + + def forward(self, x): + """Take as input the base output vector / matrix and return the diagonal covariance matrix(ces). + + Args: + x (torch.Tensor): base output vector / matrix of shape (N, B) + + Returns: + torch.Tensor: covariance matrices (one covariance matrix for each sample) of shape (N, D, D) + """ + # from base outputs to vector representing the log var + x = self._model(x) # shape (N,D) + + # get variance by taking the exponential (which is always positive) + x = torch.exp(x) # shape (N,D) + + # create covariance matrix of shape (N,D,D) + covariance = torch.zeros(x.size(0), self._dim, self._dim) # shape (N,D,D) + + # set the elements in the covariance matrix # TODO: remove this for-loop + for i in range(len(covariance)): + covariance[i] = torch.diag(x[i]) + self._offset + + # return the diagonal covariance matrix + return covariance + + +class FullCovarianceModule(torch.nn.Module): + r"""Full Covariance Generator + + This class generates a full covariance matrix given the base output vector / matrix, and maps it to the lower + triangular matrix of the Cholesky decomposition of the covariance matrix, which is then multiplied by its + transpose to give back the full covariance matrix. + + A quick reminder: + + Because the covariance is a symmetric, positive semi-definite matrix, it has a Cholesky decomposition. Thus, it + can be expressed as the product of a lower triangular matrix with its transpose :math:`\Sigma = LL^\top`. + This is useful for two reasons: + - any lower triangular matrices multiplied with its transpose results in a symmetric positive semi-definite + matrix, which thus represents a proper covariance matrix. This can be for instance useful when predicting a + full covariance matrix with a neural network. Indeed, it is hard to enforce that type of constraint (i.e. making + sure that the produced covariance matrix is symmetric and positive semi-definite) while optimizing the network. + We can thus instead output a lower-triangular matrix and multiplied by its transpose. + - it allows to solve efficiently a system of linear equations :math:`Ax = b` without having to compute the inverse + (and thus, the determinant). This is achieved in a 2-step way, by first computing :math:`Ly=b` for :math:`y` by + forward substitution, and then computing :math:`L^\top x = y` by backward substitution. + """ + + def __init__(self, num_inputs, num_outputs, offset=1.e-4): + """ + Initialize the full covariance generator. + + Args: + num_inputs (int): size of the base output vector + num_outputs (int): size of the action mean vector + offset (float): small offset to be added to the diagonal elements of the covariance matrix such that + it is positive definite. + """ + super(FullCovarianceModule, self).__init__() + # size of the mean vector + self._dim = num_outputs + + # indices for lower triangular matrix + self._idx = np.tril_indices(self._dim) + + # define diagonal offset such that the covariance is positive definite + self._offset = offset * torch.diag(torch.ones(self._dim)) + + # linear mapping between the base output vector / matrix to the covariance's lower triangular matrix + model = torch.nn.Linear(in_features=num_inputs, out_features=len(self._idx[0])) + self._model = init_orthogonal_weights_and_constant_biases(model) + + def forward(self, x): + """Take as input the base output vector / matrix and return the full covariance matrix(ces). + + Args: + x (torch.Tensor): base output vector / matrix of shape (N, B) + + Returns: + torch.Tensor: covariance matrices (one covariance matrix for each sample) of shape (N, D, D) + """ + # from base outputs to vector representing the elements of a triangular matrix + x = self._model(x) # shape (N,L) where L=(D^2+D)/2 + + # create lower triangular matrix + covariance = torch.zeros(x.size(0), self._dim, self._dim) # shape (N,D,D) + + # set the elements in the covariance matrix # TODO: find a better way than a for-loop + for i in range(len(covariance)): + covariance[i][self._idx] = x[i] + covariance[i] = torch.dot(covariance[i], covariance[i].T) + self._offset + + # add a small noise to the diagonal terms to be positive definite + # covariance = covariance + self.threshold * torch.diag(torch.ones(self.dim)) # shape (N,D,D) + + # return the full covariance matrix + return covariance + + +# define aliases for logits and probs module +FixedLogitsModule = FixedVectorModule +FixedProbsModule = FixedVectorModule +LogitsModule = VectorModule +ProbsModule = VectorModule + + +class GaussianModule(torch.nn.Module): + r"""Gaussian Module + + Type: continuous + + The Gaussian module accepts as inputs the mean and covariance modules (i.e. that inherit from `torch.nn.Modules`), + and returns the multivariate Gaussian distribution (that inherits from `torch.distributions.MultivariateNormal`). + For more information about this distribution, see the documentation of `pyrobolearn/distributions/gaussian.py`. + + Examples: + >>> # fixed gaussian (that has a fixed mean and diagonal covariance) + >>> mean = FixedMeanModule(mean=torch.zeros(5)) + >>> covariance = FixedDiagonalCovarianceModule(variances=torch.ones(5)) + >>> fixed_gaussian = GaussianModule(mean=mean, covariance=covariance) + >>> probs = fixed_gaussian(base_output) # or fixed_gaussian(output) + >>> # most flexible gaussian (that learns the mean and the full covariance) + >>> mean = MeanModule(num_inputs=10, num_outputs=5) + >>> covariance = FullCovarianceModule(num_inputs=10, num_outputs=5) + >>> full_gaussian = GaussianModule(mean=mean, covariance=covariance) + >>> probs = full_gaussian(base_output) + >>> # if the mean is already computed elsewhere + >>> mean = IdentityModule() + >>> covariance = FullCovarianceModule(num_inputs=5, num_outputs=5) + >>> gaussian = GaussianModule(mean=mean, covariance=covariance) + >>> probs = gaussian(output, base_output) # inputs for the mean and covariance + """ + + def __init__(self, mean, covariance): + """ + Initialize the Gaussian module. + + Args: + mean (torch.nn.Module): mean module. + covariance (torch.nn.Module): covariance module. + """ + super(GaussianModule, self).__init__() + self.mean = mean + self.covariance = covariance + + @property + def mean(self): + """Return the mean module.""" + return self._mean + + @mean.setter + def mean(self, mean): + """Set the mean module.""" + if not isinstance(mean, torch.nn.Module): + raise TypeError("Expecting the mean to be an instance of `torch.nn.Module`, instead got: " + "{}".format(type(mean))) + self._mean = mean + + @property + def covariance(self): + """Return the covariance module.""" + return self._covariance + + @covariance.setter + def covariance(self, covariance): + """Set the covariance module.""" + """Set the covariance.""" + if not isinstance(covariance, torch.nn.Module): + raise TypeError("Expecting the covariance to be an instance of `torch.nn.Module`, instead got: " + "{}".format(type(covariance))) + self._covariance = covariance + + def forward(self, *x): + """Forward the given inputs :attr:`x`.""" + if len(x) == 1: + x1, x2 = x[0], x[0] + elif len(x) == 2: + x1, x2 = x[0], x[1] + else: + raise ValueError("Expecting 1 or 2 inputs.") + mean = self.mean(x1) + covariance = self.covariance(x2) + return GaussianDistribution(mean=mean, covariance=covariance) + + +class DiscreteModule(torch.nn.Module): + r"""Discrete probability module. + + Discrete probability module from which several discrete probability distributions (such as Bernoulli, Categorical, + and others) inherit from. + """ + + __metaclass__ = ABCMeta + + def __init__(self, probs=None, logits=None): + """ + Initialize the Discrete probability module. + + Args: + probs (torch.nn.Module): event probabilities module. + logits (torch.nn.Module): event logits module. + """ + super(DiscreteModule, self).__init__() + self.logits = logits + self.probs = probs + + @property + def logits(self): + """Return the logits module.""" + return self._logits + + @logits.setter + def logits(self, logits): + """Set the logits module.""" + if logits is not None and not isinstance(logits, torch.nn.Module): + raise TypeError("Expecting the logits to be an instance of `torch.nn.Module`, instead got: " + "{}".format(type(logits))) + self._logits = logits + self._probs = lambda x: None + + @property + def probs(self): + """Return the probabilities module.""" + return self._probs + + @probs.setter + def probs(self, probs): + """Set the probabilities module.""" + if probs is not None and not isinstance(probs, torch.nn.Module): + raise TypeError("Expecting the probs to be an instance of `torch.nn.Module`, instead got: " + "{}".format(type(probs))) + self._probs = probs + self._logits = lambda x: None + + +class CategoricalModule(DiscreteModule): + r"""Categorical Module + + Type: discrete, multiple categories + + The Categorical module accepts as inputs the discrete logits or probabilities modules, and returns the categorical + distribution (that inherits from `torch.distributions.Categorical`). + + Description: "A categorical distribution (also called a generalized Bernoulli distribution, multinoulli + distribution) is a discrete probability distribution that describes the possible results of a random variable that + can take on one of K possible categories, with the probability of each category separately specified." [1] + + Examples: + >>> # fixed categorical + >>> logits = FixedLogitsModule(torch.ones(5)) + >>> categorical = CategoricalModule(logits=logits) + >>> probs = categorical(base_output) # or categorical(output) + >>> # flexible categorical + >>> logits = LogitsModule(num_inputs=10, num_outputs=5) + >>> categorical = CategoricalModule(logits=logits) + >>> probs = categorical(base_output) + >>> # identity categorical (just copy what will be given as inputs) + >>> logits = IdentityModule() + >>> categorical = CategoricalModule(logits=logits) + >>> probs = categorical(output) + + References: + [1] Categorical distribution: https://en.wikipedia.org/wiki/Categorical_distribution + """ + + def __init__(self, probs=None, logits=None): + """ + Initialize the Categorical module. + + Args: + probs (torch.nn.Module): event probabilities module. + logits (torch.nn.Module): event logits module. + """ + super(CategoricalModule, self).__init__(probs=probs, logits=logits) + + def forward(self, x): + """Forward the given inputs :attr:`x`.""" + return CategoricalDistribution(probs=self.probs(x), logits=self.logits(x)) + + +class BernoulliModule(DiscreteModule): + r"""Bernoulli Module + + Type: discrete, binary + + "The Bernoulli distribution is the discrete probability distribution of a random variable which takes the value 1 + with probability :math:`p` and the value 0 with probability :math:`q = 1-p`, that is, the probability distribution + of any single experiment that asks a yes/no question; the question results in a boolean-valued outcome, a single + bit of information whose value is success with probability :math:`p` and failure with probability :math:`q`." [1] + + Examples: + >>> # fixed categorical + >>> logits = FixedLogitsModule(torch.ones(5)) + >>> bernoulli = BernoulliModule(logits=logits) + >>> probs = bernoulli(base_output) # or bernoulli(output) + >>> # flexible categorical + >>> logits = LogitsModule(num_inputs=10, num_outputs=5) + >>> bernoulli = BernoulliModule(logits=logits) + >>> probs = bernoulli(base_output) + >>> # identity categorical (just copy what will be given as inputs) + >>> logits = IdentityModule() + >>> bernoulli = BernoulliModule(logits=logits) + >>> probs = bernoulli(output) + + References: + [1] Bernoulli distribution: https://en.wikipedia.org/wiki/Bernoulli_distribution + """ + + def __init__(self, probs=None, logits=None): + """ + Initialize the Bernoulli module. + + Args: + probs (torch.nn.Module): event probabilities module. + logits (torch.nn.Module): event logits module. + """ + super(BernoulliModule, self).__init__(probs=probs, logits=logits) + self.logits = logits + self.probs = probs + + def forward(self, x): + """Forward the given inputs :attr:`x`.""" + return BernoulliDistribution(probs=self.probs(x), logits=self.logits(x)) diff --git a/pyrobolearn/models/gaussian.py b/pyrobolearn/models/gaussian.py index d82fb67..cca0d57 100755 --- a/pyrobolearn/models/gaussian.py +++ b/pyrobolearn/models/gaussian.py @@ -234,13 +234,13 @@ class Gaussian(object): @staticmethod def is_parametric(): """The Gaussian distribution is a nonparametric model; the mean and covariance summarized the data""" - return True + return False @staticmethod def is_linear(): """The Gaussian doesn't have parameters. Even if the mean and covariance are considered as parameters, the model is not linear wrt them""" - return True + return False @staticmethod def is_recurrent(): @@ -250,17 +250,17 @@ class Gaussian(object): @staticmethod def is_probabilistic(): """The Gaussian distribution is by definition a probabilistic model""" - return False + return True @staticmethod def is_discriminative(): """The Gaussian is not a discriminative model; no inputs are involved""" - return True + return False @staticmethod def is_generative(): """The Gaussian is a generative model, and thus we can sample from it""" - return False + return True @staticmethod def compute_mean(X, axis=0): @@ -1227,21 +1227,21 @@ MVN = Gaussian # Plotting functions # ###################### -def plot3D(ax, X, Y, pdf, title=None, xlabel='x1', ylabel='x2', zlabel='x3'): +def plot_3d(ax, X, Y, pdf, title=None, xlabel='x1', ylabel='x2', zlabel='x3'): if isinstance(pdf, (tuple, list)): pdf = np.max(np.dstack(pdf), axis=-1) ax.set(title=title, xlabel=xlabel, ylabel=ylabel, zlabel=zlabel) ax.plot_surface(X, Y, pdf, cmap='viridis', linewidth=0) -def plot2DContour(ax, x, y, pdf, title=None, xlabel='x1', ylabel='x2'): +def plot_2d_contour(ax, x, y, pdf, title=None, xlabel='x1', ylabel='x2'): if isinstance(pdf, (tuple, list)): pdf = np.max(np.dstack(pdf), axis=-1) ax.set(title=title, xlabel=xlabel, ylabel=ylabel) ax.contourf(x, y, pdf) -def plot3DAnd2DCountour(gaussians, step=500, bound=10, fig=None, title='', block=True): +def plot_3d_and_2d_countour(gaussians, step=500, bound=10, fig=None, title='', block=True): if not isinstance(gaussians, (list, tuple)): gaussians = [gaussians] @@ -1265,18 +1265,18 @@ def plot3DAnd2DCountour(gaussians, step=500, bound=10, fig=None, title='', block # 1st subplot (3D) ax = fig.add_subplot(1, 2, 1, projection='3d') - plot3D(ax, X, Y, pdf, title='p(x1, x2)', xlabel='x1', ylabel='x2', zlabel='p') + plot_3d(ax, X, Y, pdf, title='p(x1, x2)', xlabel='x1', ylabel='x2', zlabel='p') # 2nd subplot (2D) ax = fig.add_subplot(1, 2, 2) - plot2DContour(ax, x, y, pdf, title='p(x1, x2)', xlabel='x1', ylabel='x2') + plot_2d_contour(ax, x, y, pdf, title='p(x1, x2)', xlabel='x1', ylabel='x2') # show plot fig.tight_layout() plt.show(block=block) -def plot2DEllipse(ax, gaussian, color='g', fill=False, plot_2devs=False, plot_arrows=True): +def plot_2d_ellipse(ax, gaussian, color='g', fill=False, plot_2devs=False, plot_arrows=True): # alias g = gaussian @@ -1313,7 +1313,7 @@ def plot2DEllipse(ax, gaussian, color='g', fill=False, plot_2devs=False, plot_ar return ellipse_2std -def plot3DAnd2DConditional(joint_gaussian, cond_gaussian, x1_value = 0, step=500, bound=10, block=True): +def plot_3d_and_2d_conditional(joint_gaussian, cond_gaussian, x1_value=0., step=500, bound=10, block=True): # Create grid and multivariate normal x = np.linspace(-bound, bound, step) @@ -1331,7 +1331,7 @@ def plot3DAnd2DConditional(joint_gaussian, cond_gaussian, x1_value = 0, step=500 # 1st subplot (3D) ax = fig.add_subplot(1, 2, 1, projection='3d') - plot3D(ax, X, Y, joint_pdf, title='p(x1,x2)', xlabel='x1', ylabel='x2', zlabel='p') + plot_3d(ax, X, Y, joint_pdf, title='p(x1,x2)', xlabel='x1', ylabel='x2', zlabel='p') # draw plane that cut the gaussian y1 = np.linspace(-bound, bound, 2) @@ -1373,7 +1373,7 @@ if __name__ == '__main__': plt.show() # 3D and 2D plots of the Gaussian distributions - plot3DAnd2DCountour([g1, g2]) + plot_3d_and_2d_countour([g1, g2]) # Use 1 Gaussian # @@ -1402,13 +1402,13 @@ if __name__ == '__main__': fig, ax = plt.subplots(1,1) ax.set(title='Sampling from one Gaussian', aspect='equal') ax.scatter(samples[:, 0], samples[:, 1], color='b') - plot2DEllipse(ax, g2, fill=True, plot_2devs=True, plot_arrows=True) + plot_2d_ellipse(ax, g2, fill=True, plot_2devs=True, plot_arrows=True) plt.show() # conditional distribution of the Gaussian p(y|x) - x_value = 0 + x_value = 0. g_cond = g2.condition(input_value=x_value, output_idx=1) - plot3DAnd2DConditional(g2, g_cond, x_value) + plot_3d_and_2d_conditional(g2, g_cond, x_value) # marginalization of the Gaussian by summing and using the normal distribution # by summing @@ -1437,7 +1437,7 @@ if __name__ == '__main__': fig, ax = plt.subplots(1, 1) ax.set(title='Gaussian under affine transformation', aspect='equal') ax.scatter(samples[:, 0], samples[:, 1], color='b') - plot2DEllipse(ax, g_aff) + plot_2d_ellipse(ax, g_aff) plt.show() # use 2 Gaussians # @@ -1446,9 +1446,9 @@ if __name__ == '__main__': g_sum = g1 + g2 fig, ax = plt.subplots(1,1) ax.set(title='addition', xlim=[-5, 5], ylim=[-5, 5], aspect='equal') - e1 = plot2DEllipse(ax, g1, color='g', plot_arrows=False) - e2 = plot2DEllipse(ax, g2, color='b', plot_arrows=False) - e3 = plot2DEllipse(ax, g_sum, color='r', plot_arrows=False) + e1 = plot_2d_ellipse(ax, g1, color='g', plot_arrows=False) + e2 = plot_2d_ellipse(ax, g2, color='b', plot_arrows=False) + e3 = plot_2d_ellipse(ax, g_sum, color='r', plot_arrows=False) ax.legend([e1, e2, e3], ['G1', 'G2', 'G1+G2'], loc=2) plt.show() @@ -1456,9 +1456,9 @@ if __name__ == '__main__': g_mul = g1 * g2 fig, ax = plt.subplots(1, 1) ax.set(title='multiplication', xlim=[-5, 5], ylim=[-5, 5], aspect='equal') - e1 = plot2DEllipse(ax, g1, color='g', plot_arrows=False) - e2 = plot2DEllipse(ax, g2, color='b', plot_arrows=False) - e3 = plot2DEllipse(ax, g_mul, color='r', plot_arrows=False) + e1 = plot_2d_ellipse(ax, g1, color='g', plot_arrows=False) + e2 = plot_2d_ellipse(ax, g2, color='b', plot_arrows=False) + e3 = plot_2d_ellipse(ax, g_mul, color='r', plot_arrows=False) ax.legend([e1, e2, e3], ['G1', 'G2', 'G1*G2'], loc=2) plt.show() @@ -1471,8 +1471,8 @@ if __name__ == '__main__': # fit one Gaussian and plot it along the data g = Gaussian() g.fit(samples) - plot3DAnd2DCountour(g_data, title='Gaussian that generated the data', block=False) - plot3DAnd2DCountour(g, title='fitted Gaussian') + plot_3d_and_2d_countour(g_data, title='Gaussian that generated the data', block=False) + plot_3d_and_2d_countour(g, title='fitted Gaussian') # TODO # fit Gaussian on different manifolds #