From 2d537734063fba796123316796c69cd3bb5deb98 Mon Sep 17 00:00:00 2001 From: Brian Delhaisse Date: Fri, 19 Apr 2019 04:42:25 +0200 Subject: [PATCH] update distributions + fix errors in losses/explorations --- pyrobolearn/algos/evaluator.py | 2 +- pyrobolearn/algos/explorer.py | 4 +- pyrobolearn/algos/ppo.py | 4 +- pyrobolearn/distributions/bernoulli.py | 32 +++++++ pyrobolearn/distributions/categorical.py | 32 +++++++ pyrobolearn/distributions/gaussian.py | 94 +++++++++++++------ pyrobolearn/exploration/actions/boltzmann.py | 6 +- pyrobolearn/exploration/actions/eps_greedy.py | 15 ++- pyrobolearn/exploration/actions/gaussian.py | 4 +- pyrobolearn/losses/policy_losses.py | 66 ++++++++----- pyrobolearn/storages/storage.py | 16 ++++ 11 files changed, 207 insertions(+), 68 deletions(-) diff --git a/pyrobolearn/algos/evaluator.py b/pyrobolearn/algos/evaluator.py index 9ee0c24..edfd5c3 100644 --- a/pyrobolearn/algos/evaluator.py +++ b/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 ####") ############# diff --git a/pyrobolearn/algos/explorer.py b/pyrobolearn/algos/explorer.py index eedb38e..c9cbb66 100644 --- a/pyrobolearn/algos/explorer.py +++ b/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'])) diff --git a/pyrobolearn/algos/ppo.py b/pyrobolearn/algos/ppo.py index c438d28..2aa4620 100755 --- a/pyrobolearn/algos/ppo.py +++ b/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') diff --git a/pyrobolearn/distributions/bernoulli.py b/pyrobolearn/distributions/bernoulli.py index ef5e4e9..2f20532 100644 --- a/pyrobolearn/distributions/bernoulli.py +++ b/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]) diff --git a/pyrobolearn/distributions/categorical.py b/pyrobolearn/distributions/categorical.py index d767db3..39d82b6 100644 --- a/pyrobolearn/distributions/categorical.py +++ b/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]) diff --git a/pyrobolearn/distributions/gaussian.py b/pyrobolearn/distributions/gaussian.py index cda7ba1..ab0066f 100644 --- a/pyrobolearn/distributions/gaussian.py +++ b/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): """ diff --git a/pyrobolearn/exploration/actions/boltzmann.py b/pyrobolearn/exploration/actions/boltzmann.py index 26290fe..7c0a61c 100644 --- a/pyrobolearn/exploration/actions/boltzmann.py +++ b/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 diff --git a/pyrobolearn/exploration/actions/eps_greedy.py b/pyrobolearn/exploration/actions/eps_greedy.py index e086e2e..011562a 100644 --- a/pyrobolearn/exploration/actions/eps_greedy.py +++ b/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 diff --git a/pyrobolearn/exploration/actions/gaussian.py b/pyrobolearn/exploration/actions/gaussian.py index 668eb07..29fb61e 100644 --- a/pyrobolearn/exploration/actions/gaussian.py +++ b/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 diff --git a/pyrobolearn/losses/policy_losses.py b/pyrobolearn/losses/policy_losses.py index 366dfd8..d3673b4 100644 --- a/pyrobolearn/losses/policy_losses.py +++ b/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 diff --git a/pyrobolearn/storages/storage.py b/pyrobolearn/storages/storage.py index 5b8c55b..d6b5fd1 100644 --- a/pyrobolearn/storages/storage.py +++ b/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()