update distribution modules/layers

.gitignore

@@ -109,4 +109,3 @@ venv.bak/
 .ipynb_checkpoints/
 .idea/
 0_VRDemoSettings.txt
-tests/

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.

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))

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.
+```

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 *

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()

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)

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

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))

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 #