update distributions + fix errors in losses/explorations

pyrobolearn/algos/evaluator.py

@@ -86,7 +86,7 @@ class Evaluator(object):
             returns = self.estimator.evaluate(self.storage)
 
             if verbose:
-                print("Returns: {}".format(returns))
+                # print("Returns: {}".format(returns))
                 print("#### End of the Evaluation phase ####")
 
     #############

pyrobolearn/algos/explorer.py

@@ -173,9 +173,11 @@ class Explorer(object):
 
             # if done, get out of the loop
             if done:
-                self.storage.end(rollout_idx)  # fill remaining mask values
                 break
 
+        # fill remaining mask values
+        self.storage.end(rollout_idx)
+
         if verbose:
             print("#### End of the Exploration phase #####")
             # print("states: {}".format(self.storage['states']))

pyrobolearn/algos/ppo.py

@@ -17,7 +17,7 @@ from pyrobolearn.exploration import ActionExploration
 from pyrobolearn.storages import RolloutStorage
 from pyrobolearn.samplers import BatchRandomSampler
 from pyrobolearn.returns import GAE, PolicyEvaluator
-from pyrobolearn.losses import CLIPLoss, L2Loss, EntropyLoss
+from pyrobolearn.losses import CLIPLoss, ValueL2Loss, EntropyLoss
 from pyrobolearn.optimizers import Adam
 
 from pyrobolearn import logger
@@ -198,7 +198,7 @@ class PPO(GradientRLAlgo):
 
         # create loss
         logger.debug('create loss')
-        loss = CLIPLoss(estimator, clip=clip) + l2_coeff * L2Loss(estimator, value) + entropy_coeff * EntropyLoss()
+        loss = CLIPLoss(estimator, clip=clip) + l2_coeff * ValueL2Loss(estimator, value) + entropy_coeff * EntropyLoss()
 
         # create optimizer
         logger.debug('create Adam optimizer')

pyrobolearn/distributions/bernoulli.py

@@ -44,3 +44,35 @@ class Bernoulli(torch.distributions.Bernoulli):
     def mode(self):
         """Return the mode of the Bernoulli distribution."""
         return torch.gt(self.probs, 0.5).float()
+
+    @staticmethod
+    def from_list(bernoullis):
+        """
+        Convert a list of Bernoulli [Ber1, Ber2, ..., BerN] to a single Bernoulli distribution with N logits /
+        probs.
+
+        Args:
+            bernoullis (list of Bernoulli): list of Bernoulli distributions.
+
+        Returns:
+            Bernoulli: resulting single Bernoulli distribution.
+        """
+        return Bernoulli(probs=torch.stack([bernoulli.probs for bernoulli in bernoullis]))
+
+    def __getitem__(self, indices):
+        """
+        Get the corresponding Bernoullis from the single Bernoulli distribution. That is, if the single Bernoulli
+        distribution has multiple logits / probs, it selects the corresponding Bernoulli distributions from it.
+
+        Examples:
+            bernoulli = Bernoulli(probs=torch.tensor([[0.25, 0.75], [0.6, 0.4], [0.7, 0.3]]))
+            bernoulli[[0,2]]  # this returns Bernoulli(probs=torch.tensor([[0.25, 0.75], [0.7, 0.3]]))
+            bernoulli[:2]  # this returns Bernoulli(probs=torch.tensor([[0.25, 0.75], [0.6, 0.4]]))
+
+        Args:
+            indices (int, list of int, slices): indices.
+
+        Returns:
+            Bernoulli: resulting sliced Bernoulli distribution.
+        """
+        return Bernoulli(probs=self.probs[indices])

pyrobolearn/distributions/categorical.py

@@ -45,3 +45,35 @@ class Categorical(torch.distributions.Categorical):
     def mode(self):
         """Return the mode of the Categorical distribution."""
         return self.probs.argmax(dim=-1, keepdim=True)
+
+    @staticmethod
+    def from_list(categoricals):
+        """
+        Convert a list of Categorical [cat1, cat2, ..., catN] to a single Categorical distribution with N logits /
+        probs.
+
+        Args:
+            categoricals (list of Categorical): list of Categorical distributions.
+
+        Returns:
+            Categorical: resulting single Categorical distribution.
+        """
+        return Categorical(probs=torch.stack([categorical.probs for categorical in categoricals]))
+
+    def __getitem__(self, indices):
+        """
+        Get the corresponding Categoricals from the single Categorical distribution. That is, if the single Categorical
+        distribution has multiple logits / probs, it selects the corresponding Categorical distributions from it.
+
+        Examples:
+            categorical = Categorical(probs=torch.tensor([[0.25, 0.75], [0.6, 0.4], [0.7, 0.3]]))
+            categorical[[0,2]]  # this returns Categorical(probs=torch.tensor([[0.25, 0.75], [0.7, 0.3]]))
+            categorical[:2]  # this returns Categorical(probs=torch.tensor([[0.25, 0.75], [0.6, 0.4]]))
+
+        Args:
+            indices (int, list of int, slices): indices.
+
+        Returns:
+            Categorical: resulting sliced Categorical distribution.
+        """
+        return Categorical(probs=self.probs[indices])

pyrobolearn/distributions/gaussian.py

@@ -884,6 +884,22 @@ class Gaussian(torch.distributions.MultivariateNormal):
     # def __array_ufunc__(self, *args):
     #     print(args)
 
+    @staticmethod
+    def from_list(gaussians):
+        """
+        Convert a list of Gaussian [gauss1, gauss2, ..., gaussN] to a single Gaussian distribution with N means and
+        covariances.
+
+        Args:
+            gaussians (list of Gaussian): list of Gaussian distributions.
+
+        Returns:
+            Gaussian: resulting single Gaussian distribution.
+        """
+        means = torch.stack([gaussian.mean for gaussian in gaussians])
+        covariances = torch.stack([gaussian.covariance for gaussian in gaussians])
+        return Gaussian(mean=means, covariance=covariances)
+
     #############
     # Operators #
     #############
@@ -909,41 +925,61 @@ class Gaussian(torch.distributions.MultivariateNormal):
             return self.pdf(x)
         return self.sample_n(n=size)
 
-    def __getitem__(self, idx):
+    def __getitem__(self, indices):
         """
-        Conditional and marginal distribution.
-
-        Args:
-            idx (int, slice, tuple): if int or slice, it will return the marginal distribution. If tuple, it
-                will return the conditional distribution.
-
-        Returns:
-            Gaussian: conditional or marginal distribution
+        Get the corresponding Gaussians from the single Gaussian distribution. That is, if the single Gaussian
+        distribution has multiple means and covariances, it selects the corresponding Gaussian distributions from it.
 
         Examples:
-            # joint distribution p(x1,x2)
-            g = Gaussian(torch.Tensor([1.,2.]), np.identity(2))
+            means = torch.tensor([[0., 0.], [1., 1.], [2., 2.]])
+            covariances = torch.stack([torch.eye(2), 0.5*torch.eye(2), 2*torch.eye(2)])
+            gaussian = Gaussian(mean=means, covariance=covariances)
+            gaussian[[0,2]]  # this returns Gaussian(mean=means[[0,2]], covariance=covariances[[0,2]])
+            gaussian[:2]  # this returns Gaussian(mean=means[:2], covariance=covariances[:2])
 
-            # marginal distribution p(x1) and p(x2)
-            marg1 = g[0]
-            marg2 = g[1]
+        Args:
+            indices (int, list of int, slices): indices.
 
-            # conditional distribution p(x1|x2) and p(x2|x1)
-            sample = g.sample()
-            cond1 = g[0,1,sample[0]]
-            cond2 = g[1,0,sample[1]]
+        Returns:
+            Bernoulli: resulting sliced Bernoulli distribution.
         """
-        if isinstance(idx, tuple):  # conditional distribution
-            if len(idx) == 2:
-                value, idx1 = idx
-                idx2 = None
-            elif len(idx) == 3:
-                value, idx1, idx2 = idx
-            else:
-                raise IndexError("Expecting two or three indices: value, idx1 (, idx2)")
-            return self.condition(value, idx1, idx2)
-        else:  # marginal distribution
-            return self.marginalize(idx)
+        return Gaussian(mean=self.mean[indices], covariance=self.covariance[indices])
+
+    # def __getitem__(self, idx):
+    #     """
+    #     Conditional and marginal distribution.
+    #
+    #     Args:
+    #         idx (int, slice, tuple): if int or slice, it will return the marginal distribution. If tuple, it
+    #             will return the conditional distribution.
+    #
+    #     Returns:
+    #         Gaussian: conditional or marginal distribution
+    #
+    #     Examples:
+    #         # joint distribution p(x1,x2)
+    #         g = Gaussian(torch.Tensor([1.,2.]), np.identity(2))
+    #
+    #         # marginal distribution p(x1) and p(x2)
+    #         marg1 = g[0]
+    #         marg2 = g[1]
+    #
+    #         # conditional distribution p(x1|x2) and p(x2|x1)
+    #         sample = g.sample()
+    #         cond1 = g[0,1,sample[0]]
+    #         cond2 = g[1,0,sample[1]]
+    #     """
+    #     if isinstance(idx, tuple):  # conditional distribution
+    #         if len(idx) == 2:
+    #             value, idx1 = idx
+    #             idx2 = None
+    #         elif len(idx) == 3:
+    #             value, idx1, idx2 = idx
+    #         else:
+    #             raise IndexError("Expecting two or three indices: value, idx1 (, idx2)")
+    #         return self.condition(value, idx1, idx2)
+    #     else:  # marginal distribution
+    #         return self.marginalize(idx)
 
     def __add__(self, other):
         """

pyrobolearn/exploration/actions/boltzmann.py

@@ -61,6 +61,8 @@ class BoltzmannActionExploration(DiscreteActionExploration):
             torch.Tensor: action
             torch.distributions.Distribution: distribution on the action :math:`\pi_{\theta}(.|s)`
         """
-        distribution = self._module(outputs)
-        action = distribution.sample((1,))
+        distribution = self._module(outputs)   # shape = (N, D)
+        action = distribution.sample((1,))     # shape = (1, N)
+        if len(action.shape) > 1:
+            action = action.view(-1, 1)        # shape = (N, 1) otherwise shape = (1,)
         return action, distribution

pyrobolearn/exploration/actions/eps_greedy.py

@@ -54,9 +54,14 @@ class EpsilonGreedyActionExploration(DiscreteActionExploration):
             torch.Tensor: action
             torch.distributions.Distribution: distribution on the action :math:`\pi_{\theta}(.|s)`
         """
-        idx = torch.argmax(outputs)
-        probs = self.epsilon/(outputs.size()[-1] - 1) * torch.ones_like(outputs)
-        probs[idx] = (1. - self.epsilon)
-        distribution = Categorical(probs=probs)
-        action = distribution.sample((1,))
+        idx = torch.argmax(outputs, dim=-1)  # shape = (N,) or (1,)
+        probs = self.epsilon/(outputs.size()[-1] - 1) * torch.ones_like(outputs)  # shape = (N, D) or (D,)
+        if len(probs.shape) > 1:  # multiple data
+            probs[range(len(outputs)), idx] = (1. - self.epsilon)
+        else:
+            probs[idx] = (1. - self.epsilon)
+        distribution = Categorical(probs=probs)   # shape = (N, D)
+        action = distribution.sample((1,))        # shape = (1, N)
+        if len(action.shape) > 1:  # multiple data
+            action = action.view(-1, 1)           # shape = (N, 1) otherwise shape = (1,)
         return action, distribution

pyrobolearn/exploration/actions/gaussian.py

@@ -65,6 +65,6 @@ class GaussianActionExploration(ContinuousActionExploration):
             torch.Tensor: action
             torch.distributions.Distribution: distribution on the action :math:`\pi_{\theta}(.|s)`
         """
-        distribution = self._module(outputs)
-        action = distribution.rsample((1,))
+        distribution = self._module(outputs)    # shape = (N, D) or (D,)
+        action = distribution.rsample((1,))[0]  # shape = (N, D) or (D,)
         return action, distribution

pyrobolearn/losses/policy_losses.py

@@ -52,7 +52,7 @@ class PGLoss(BatchLoss):
                             "{}".format(type(estimator)))
         self._estimator = estimator
 
-    def _compute(self, batch):
+    def _compute(self, batch):  # TODO: think when there are multiple actions --> independent joint distribution?
         """
         Compute the PG loss on the given batch.
 
@@ -62,7 +62,7 @@ class PGLoss(BatchLoss):
         Returns:
             torch.Tensor: loss scalar value
         """
-        # evaluate the action  # TODO: think when there are multiple actions --> independent joint distribution?
+        # evaluate the action
         log_curr_pi = 0
         for pi, action in zip(batch.current['action_distributions'], batch.current['actions']):
             log_curr_pi += pi.log_prob(action)
@@ -106,7 +106,7 @@ class CPILoss(BatchLoss):
                             "{}".format(type(estimator)))
         self._estimator = estimator
 
-    def _compute(self, batch):
+    def _compute(self, batch):  # TODO: think when there are multiple actions
         """
         Compute the CPI loss on the given batch.
 
@@ -116,17 +116,23 @@ class CPILoss(BatchLoss):
         Returns:
             torch.Tensor: loss scalar value
         """
-        # evaluate the actions # TODO: think when there are multiple actions
-        log_curr_pi = 0
-        for pi, action in zip(batch.current['action_distributions'], batch.current['actions']):
-            log_curr_pi += pi.log_prob(action)
+        masks = batch['masks']  # shape = (len(batch), 1)
+        idx = batch.indices  # shape = (I,)  where 0 <= I <= len(batch)
 
+        # evaluate the old actions using the old distribution
         log_prev_pi = 0
         for pi, action in zip(batch['action_distributions'], batch['actions']):
-            log_prev_pi += pi.log_prob(action)
+            log_prev_pi += pi.log_prob(action) * masks  # shape = (len(batch), 1)
+        log_prev_pi = log_prev_pi[idx]  # shape = (I, 1)
 
-        ratio = torch.exp(log_curr_pi - log_prev_pi)
-        estimator = batch[self._estimator]
+        # evaluate the old actions using the current distribution
+        log_curr_pi = 0
+        for pi, action in zip(batch.current['action_distributions'], batch['actions']):
+            log_curr_pi += pi.log_prob(action) * masks  # shape = (len(batch), 1)
+        log_curr_pi = log_curr_pi[idx]  # shape = (I, 1)
+
+        ratio = torch.exp(log_curr_pi - log_prev_pi)  # shape = (I, 1)
+        estimator = batch[self._estimator][idx]  # shape = (I, 1)
 
         loss = ratio * estimator
         return -loss.mean()
@@ -166,7 +172,7 @@ class CLIPLoss(BatchLoss):
                             "{}".format(type(estimator)))
         self._estimator = estimator
 
-    def _compute(self, batch):
+    def _compute(self, batch):  # TODO: think when there are multiple actions
         """
         Compute the CLIP loss on the given batch.
 
@@ -176,19 +182,25 @@ class CLIPLoss(BatchLoss):
         Returns:
             torch.Tensor: loss scalar value
         """
-        # evaluate the actions # TODO: think when there are multiple actions
-        log_curr_pi = 0
-        for pi, action in zip(batch.current['action_distributions'], batch.current['actions']):
-            log_curr_pi += pi.log_prob(action)
+        masks = batch['masks']  # shape = (len(batch), 1)
+        idx = batch.indices     # shape = (I,)  where 0 <= I <= len(batch)
 
+        # evaluate the old actions using the old distribution
         log_prev_pi = 0
         for pi, action in zip(batch['action_distributions'], batch['actions']):
-            log_prev_pi += pi.log_prob(action)
+            log_prev_pi += pi.log_prob(action) * masks  # shape = (len(batch), 1)
+        log_prev_pi = log_prev_pi[idx]  # shape = (I, 1)
 
-        ratio = torch.exp(log_curr_pi - log_prev_pi)
-        estimator = batch[self._estimator]
+        # evaluate the old actions using the current distribution
+        log_curr_pi = 0
+        for pi, action in zip(batch.current['action_distributions'], batch['actions']):
+            log_curr_pi += pi.log_prob(action) * masks  # shape = (len(batch), 1)
+        log_curr_pi = log_curr_pi[idx]  # shape = (I, 1)
 
-        loss = torch.min(ratio * estimator, torch.clamp(ratio, 1.0-self.eps, 1.0+self.eps) * estimator)
+        ratio = torch.exp(log_curr_pi - log_prev_pi)  # shape = (I, 1)
+        estimator = batch[self._estimator][idx]       # shape = (I, 1)
+
+        loss = torch.min(ratio * estimator, torch.clamp(ratio, 1.0 - self.eps, 1.0 + self.eps) * estimator)
         return -loss.mean()
 
     def latex(self):
@@ -201,7 +213,7 @@ class KLPenaltyLoss(BatchLoss):
 
     KL Penalty to minimize:
 
-    .. math:: L^{KL}(\theta) = \mathbb{E}[ KL( \pi_{\theta_{old}}(a_t | s_t) || \pi_{\theta}(a_t | s_t) ) ]
+    .. math:: L^{KL}(\theta) = \mathbb{E}[ KL( \pi_{\theta_{old}}(\cdot | s_t) || \pi_{\theta}(\cdot | s_t) ) ]
 
     where :math:`KL(.||.)` is the KL-divergence between two probability distributions.
     """
@@ -211,12 +223,10 @@ class KLPenaltyLoss(BatchLoss):
         Initialize the KL Penalty loss.
         """
         super(KLPenaltyLoss, self).__init__()
-        # self.p = p
-        # self.q = q
 
     def _compute(self, batch):
-        """
-        Compute the KL divergence loss: :math:`KL(p||q)`.
+        r"""
+        Compute the KL divergence loss: :math:`KL[\pi_{\theta_{old}}(\cdot | s_t) || \pi_{\theta}(\cdot | s_t)]`.
 
         Args:
             batch (Batch): batch containing the states, actions, rewards, etc.
@@ -224,9 +234,11 @@ class KLPenaltyLoss(BatchLoss):
         Returns:
             torch.Tensor: loss scalar value
         """
+        idx = batch.indices  # shape = (I,)  where 0 <= I <= len(batch)
+
         kl_div = 0
         for curr_pi, prev_pi in zip(batch.current['action_distributions'], batch['action_distributions']):
-            kl_div += torch.distributions.kl.kl_divergence(prev_pi, curr_pi).mean()
+            kl_div += torch.distributions.kl.kl_divergence(prev_pi[idx], curr_pi[idx]).mean()
         return kl_div
 
     def latex(self):
@@ -265,7 +277,9 @@ class EntropyLoss(BatchLoss):
         Returns:
             torch.Tensor: loss scalar value
         """
+        idx = batch.indices  # shape = (I,)  where 0 <= I <= len(batch)
+
         entropy = 0
         for dist in batch.current['action_distributions']:
-            entropy += dist.entropy().mean()
+            entropy += dist[idx].entropy().mean()
         return entropy

pyrobolearn/storages/storage.py

@@ -1033,6 +1033,22 @@ class RolloutStorage(DictStorage):  # TODO: think about when multiple policies:
             else:  # value = tensor
                 batch[key] = sample(value, indices)
 
+        # TODO: add the following lines in the `end` method? Need to be called only one time as long we don't clear
+        # replace the zeros in the action_distributions field by dummy distributions
+        # take the first distribution for each action distribution
+        dists = [dist[dist != 0][0] for dist in self['action_distributions']]
+        # replace the zeros by that first distribution
+        for distribution, dist in zip(batch['action_distributions'], dists):
+            distribution[distribution == 0] = dist
+
+        # Now that we have a list of distributions for each action, transform it to a single distribution
+        distributions = batch['action_distributions']
+        for i, distribution in enumerate(distributions):
+            # select first distribution
+            dist = distribution[0]
+            # transform list of distributions to a single distribution and put it in the batch
+            distributions[i] = dist.__class__.from_list(distribution)
+
         # create Batch object
         batch = Batch(batch, device=self.device, dtype=self.dtype)
         batch.indices = torch.tensor(range(len(indices)))[batch['masks'][:, 0] != 0].tolist()