From eca868955dd0c423f59006907efa914a94056262 Mon Sep 17 00:00:00 2001 From: OrangeX4 <318483724@qq.com> Date: Wed, 4 Sep 2024 16:08:57 +0800 Subject: [PATCH 1/2] update: add gumbel_softmax --- scripts/configs/hpl/discrete/gym.yaml | 2 + scripts/configs/hpl/discrete/metaworld.yaml | 2 + wiserl/algorithm/hpl/hpl.py | 19 ++- wiserl/utils/distributions.py | 150 +++++++++++++++++++- 4 files changed, 167 insertions(+), 6 deletions(-) diff --git a/scripts/configs/hpl/discrete/gym.yaml b/scripts/configs/hpl/discrete/gym.yaml index a9fcf1d..262099c 100644 --- a/scripts/configs/hpl/discrete/gym.yaml +++ b/scripts/configs/hpl/discrete/gym.yaml @@ -16,6 +16,8 @@ algorithm: reg_coef: 0.0 discrete: true discrete_group: 8 + gumbel_softmax: true + temperature: 1.0 stoc_encoding: false # for hopper-medium-expert, try true rm_label: true diff --git a/scripts/configs/hpl/discrete/metaworld.yaml b/scripts/configs/hpl/discrete/metaworld.yaml index 3c60ed9..61a678b 100644 --- a/scripts/configs/hpl/discrete/metaworld.yaml +++ b/scripts/configs/hpl/discrete/metaworld.yaml @@ -16,6 +16,8 @@ algorithm: reg_coef: 0.0001 discrete: true discrete_group: 8 + gumbel_softmax: true + temperature: 1.0 stoc_encoding: true rm_label: true diff --git a/wiserl/algorithm/hpl/hpl.py b/wiserl/algorithm/hpl/hpl.py index 557efd7..c04617e 100644 --- a/wiserl/algorithm/hpl/hpl.py +++ b/wiserl/algorithm/hpl/hpl.py @@ -16,6 +16,7 @@ from wiserl.module.net.mlp import MLP from wiserl.utils.functional import expectile_regression from wiserl.utils.misc import make_target, sync_target +from wiserl.utils.distributions import RelaxedOneHotCategorical class Decoder(nn.Module): @@ -70,6 +71,8 @@ def __init__( stoc_encoding: bool = True, discrete: bool = True, discrete_group: int = 8, + gumbel_softmax: bool = False, + temperature: float = 1.0, **kwargs ): self.expectile = expectile @@ -89,6 +92,8 @@ def __init__( self.stoc_encoding = stoc_encoding self.discrete = discrete self.discrete_group = discrete_group + self.gumbel_softmax = gumbel_softmax + self.temperature = temperature self.rm_label = rm_label super().__init__(*args, **kwargs) # define the attention mask for future prediction @@ -200,10 +205,16 @@ def select_reward(self, batch, deterministic=False): def get_z_distribution(self, logits): if self.discrete: logits = logits.reshape(*logits.shape[:-1], self.discrete_group, -1) - return torch.distributions.Independent( - torch.distributions.OneHotCategoricalStraightThrough(logits=logits), - reinterpreted_batch_ndims=1 - ) + if self.gumbel_softmax: + return torch.distributions.Independent( + RelaxedOneHotCategorical(self.temperature, logits=logits), + reinterpreted_batch_ndims=1 + ) + else: + return torch.distributions.Independent( + torch.distributions.OneHotCategoricalStraightThrough(logits=logits), + reinterpreted_batch_ndims=1 + ) else: mean, logstd = logits.chunk(2, dim=-1) return torch.distributions.Independent( diff --git a/wiserl/utils/distributions.py b/wiserl/utils/distributions.py index 141ecad..c927751 100644 --- a/wiserl/utils/distributions.py +++ b/wiserl/utils/distributions.py @@ -3,8 +3,10 @@ import numpy as np import torch -from torch.distributions import Normal - +from torch.distributions import Normal, constraints, OneHotCategorical +from torch.distributions.transformed_distribution import TransformedDistribution +from torch.distributions.transforms import ExpTransform +from torch.distributions.utils import broadcast_all, clamp_probs class TanhNormal(Normal): def __init__(self, @@ -40,3 +42,147 @@ def entropy(self): @property def tanh_mean(self): return torch.tanh(self.mean) + + +# Code of RelaxedOneHotCategorical came from: +# +# ``` +# from torch.distributions import RelaxedOneHotCategorical +# ``` +# +# We inherit `OneHotCategorical` to use `dist.kl_divergence()`, +# and the `torch.distributions.RelaxedOneHotCategorical` does not implement `dist.kl_divergence()`. + +class ExpRelaxedCategorical(OneHotCategorical): + r""" + Creates a ExpRelaxedCategorical parameterized by + :attr:`temperature`, and either :attr:`probs` or :attr:`logits` (but not both). + Returns the log of a point in the simplex. Based on the interface to + :class:`OneHotCategorical`. + + Implementation based on [1]. + + See also: :func:`torch.distributions.OneHotCategorical` + + Args: + temperature (Tensor): relaxation temperature + probs (Tensor): event probabilities + logits (Tensor): unnormalized log probability for each event + + [1] The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables + (Maddison et al, 2017) + + [2] Categorical Reparametrization with Gumbel-Softmax + (Jang et al, 2017) + """ + arg_constraints = {"probs": constraints.simplex, "logits": constraints.real_vector} + support = ( + constraints.real_vector + ) # The true support is actually a submanifold of this. + has_rsample = True + + def __init__(self, temperature, probs=None, logits=None, validate_args=None): + self.temperature = temperature + super().__init__(probs=probs, logits=logits, validate_args=validate_args) + + def expand(self, batch_shape, _instance=None): + new = self._get_checked_instance(ExpRelaxedCategorical, _instance) + batch_shape = torch.Size(batch_shape) + new.temperature = self.temperature + new._categorical = self._categorical.expand(batch_shape) + super(ExpRelaxedCategorical, new).__init__( + batch_shape, self.event_shape, validate_args=False + ) + new._validate_args = self._validate_args + return new + + def _new(self, *args, **kwargs): + return self._categorical._new(*args, **kwargs) + + @property + def param_shape(self): + return self._categorical.param_shape + + @property + def logits(self): + return self._categorical.logits + + @property + def probs(self): + return self._categorical.probs + + def rsample(self, sample_shape=torch.Size()): + shape = self._extended_shape(sample_shape) + uniforms = clamp_probs( + torch.rand(shape, dtype=self.logits.dtype, device=self.logits.device) + ) + gumbels = -((-(uniforms.log())).log()) + scores = (self.logits + gumbels) / self.temperature + return scores - scores.logsumexp(dim=-1, keepdim=True) + + def sample(self, sample_shape: torch.Size = torch.Size()) -> torch.Tensor: + """ + Generates a sample_shape shaped sample or sample_shape shaped batch of + samples if the distribution parameters are batched. + """ + with torch.no_grad(): + return self.rsample(sample_shape) + + def log_prob(self, value): + K = self._categorical._num_events + if self._validate_args: + self._validate_sample(value) + logits, value = broadcast_all(self.logits, value) + log_scale = torch.full_like( + self.temperature, float(K) + ).lgamma() - self.temperature.log().mul(-(K - 1)) + score = logits - value.mul(self.temperature) + score = (score - score.logsumexp(dim=-1, keepdim=True)).sum(-1) + return score + log_scale + + +class RelaxedOneHotCategorical(TransformedDistribution): + r""" + Creates a RelaxedOneHotCategorical distribution parametrized by + :attr:`temperature`, and either :attr:`probs` or :attr:`logits`. + This is a relaxed version of the :class:`OneHotCategorical` distribution, so + its samples are on simplex, and are reparametrizable. + + Example:: + + >>> # xdoctest: +IGNORE_WANT("non-deterministic") + >>> m = RelaxedOneHotCategorical(torch.tensor([2.2]), + ... torch.tensor([0.1, 0.2, 0.3, 0.4])) + >>> m.sample() + tensor([ 0.1294, 0.2324, 0.3859, 0.2523]) + + Args: + temperature (Tensor): relaxation temperature + probs (Tensor): event probabilities + logits (Tensor): unnormalized log probability for each event + """ + arg_constraints = {"probs": constraints.simplex, "logits": constraints.real_vector} + support = constraints.simplex + has_rsample = True + + def __init__(self, temperature, probs=None, logits=None, validate_args=None): + base_dist = ExpRelaxedCategorical( + temperature, probs, logits, validate_args=validate_args + ) + super().__init__(base_dist, ExpTransform(), validate_args=validate_args) + + def expand(self, batch_shape, _instance=None): + new = self._get_checked_instance(RelaxedOneHotCategorical, _instance) + return super().expand(batch_shape, _instance=new) + + @property + def temperature(self): + return self.base_dist.temperature + + @property + def logits(self): + return self.base_dist.logits + + @property + def probs(self): + return self.base_dist.probs From ea40948cd25c2e57ab8f3ba694cb9b6a9a071527 Mon Sep 17 00:00:00 2001 From: OrangeX4 Date: Fri, 6 Sep 2024 00:22:22 +0800 Subject: [PATCH 2/2] fix: update to st gumbel softmax --- scripts/configs/hpl/discrete/gym.yaml | 14 +-- wiserl/algorithm/hpl/hpl.py | 4 +- wiserl/utils/distributions.py | 149 +++----------------------- 3 files changed, 25 insertions(+), 142 deletions(-) diff --git a/scripts/configs/hpl/discrete/gym.yaml b/scripts/configs/hpl/discrete/gym.yaml index 262099c..80f77d8 100644 --- a/scripts/configs/hpl/discrete/gym.yaml +++ b/scripts/configs/hpl/discrete/gym.yaml @@ -18,7 +18,7 @@ algorithm: discrete_group: 8 gumbel_softmax: true temperature: 1.0 - stoc_encoding: false # for hopper-medium-expert, try true + stoc_encoding: false # for hopper-medium-replay, try true rm_label: true checkpoint: null @@ -31,7 +31,7 @@ wandb: entity: null project: null -env: hopper-medium-expert-v2 +env: hopper-medium-replay-v2 env_kwargs: env_wrapper: env_wrapper_kwargs: @@ -81,17 +81,17 @@ network: rm_dataset: - class: D4RLOfflineDataset - env: hopper-medium-expert-v2 + env: hopper-medium-replay-v2 batch_size: 64 # [64, 128] mode: trajectory segment_length: 100 padding_mode: none - class: IPLComparisonOfflineDataset - env: hopper-medium-expert-v2 + env: hopper-medium-replay-v2 batch_size: 8 mode: human - class: D4RLOfflineDataset - env: hopper-medium-expert-v2 + env: hopper-medium-replay-v2 batch_size: 512 mode: transition rm_dataloader: @@ -100,7 +100,7 @@ rm_dataloader: rl_dataset: - class: D4RLOfflineDataset - env: hopper-medium-expert-v2 + env: hopper-medium-replay-v2 batch_size: 512 mode: transition rl_dataloader: @@ -120,7 +120,7 @@ rm_eval: function: eval_reward_model eval_dataset_kwargs: class: IPLComparisonOfflineDataset - env: hopper-medium-expert-v2 + env: hopper-medium-replay-v2 batch_size: 32 mode: human eval: false diff --git a/wiserl/algorithm/hpl/hpl.py b/wiserl/algorithm/hpl/hpl.py index c04617e..b28828f 100644 --- a/wiserl/algorithm/hpl/hpl.py +++ b/wiserl/algorithm/hpl/hpl.py @@ -16,7 +16,7 @@ from wiserl.module.net.mlp import MLP from wiserl.utils.functional import expectile_regression from wiserl.utils.misc import make_target, sync_target -from wiserl.utils.distributions import RelaxedOneHotCategorical +from wiserl.utils.distributions import OneHotCategoricalSTGumbelSoftmax class Decoder(nn.Module): @@ -207,7 +207,7 @@ def get_z_distribution(self, logits): logits = logits.reshape(*logits.shape[:-1], self.discrete_group, -1) if self.gumbel_softmax: return torch.distributions.Independent( - RelaxedOneHotCategorical(self.temperature, logits=logits), + OneHotCategoricalSTGumbelSoftmax(self.temperature, logits=logits), reinterpreted_batch_ndims=1 ) else: diff --git a/wiserl/utils/distributions.py b/wiserl/utils/distributions.py index c927751..dea4c41 100644 --- a/wiserl/utils/distributions.py +++ b/wiserl/utils/distributions.py @@ -3,10 +3,8 @@ import numpy as np import torch -from torch.distributions import Normal, constraints, OneHotCategorical -from torch.distributions.transformed_distribution import TransformedDistribution -from torch.distributions.transforms import ExpTransform -from torch.distributions.utils import broadcast_all, clamp_probs +from torch.distributions import Normal, OneHotCategorical +from torch.distributions.utils import clamp_probs class TanhNormal(Normal): def __init__(self, @@ -44,145 +42,30 @@ def tanh_mean(self): return torch.tanh(self.mean) -# Code of RelaxedOneHotCategorical came from: -# -# ``` -# from torch.distributions import RelaxedOneHotCategorical -# ``` -# -# We inherit `OneHotCategorical` to use `dist.kl_divergence()`, -# and the `torch.distributions.RelaxedOneHotCategorical` does not implement `dist.kl_divergence()`. - -class ExpRelaxedCategorical(OneHotCategorical): +class OneHotCategoricalSTGumbelSoftmax(OneHotCategorical): r""" - Creates a ExpRelaxedCategorical parameterized by - :attr:`temperature`, and either :attr:`probs` or :attr:`logits` (but not both). - Returns the log of a point in the simplex. Based on the interface to - :class:`OneHotCategorical`. - - Implementation based on [1]. - - See also: :func:`torch.distributions.OneHotCategorical` - - Args: - temperature (Tensor): relaxation temperature - probs (Tensor): event probabilities - logits (Tensor): unnormalized log probability for each event - - [1] The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables - (Maddison et al, 2017) + Creates a reparameterizable :class:`OneHotCategorical` distribution based on the straight- + through gumbel-softmax estimator from [1]. - [2] Categorical Reparametrization with Gumbel-Softmax + [1] Categorical Reparametrization with Gumbel-Softmax (Jang et al, 2017) """ - arg_constraints = {"probs": constraints.simplex, "logits": constraints.real_vector} - support = ( - constraints.real_vector - ) # The true support is actually a submanifold of this. has_rsample = True - + def __init__(self, temperature, probs=None, logits=None, validate_args=None): + super().__init__(probs, logits, validate_args=validate_args) self.temperature = temperature - super().__init__(probs=probs, logits=logits, validate_args=validate_args) - - def expand(self, batch_shape, _instance=None): - new = self._get_checked_instance(ExpRelaxedCategorical, _instance) - batch_shape = torch.Size(batch_shape) - new.temperature = self.temperature - new._categorical = self._categorical.expand(batch_shape) - super(ExpRelaxedCategorical, new).__init__( - batch_shape, self.event_shape, validate_args=False - ) - new._validate_args = self._validate_args - return new - - def _new(self, *args, **kwargs): - return self._categorical._new(*args, **kwargs) - - @property - def param_shape(self): - return self._categorical.param_shape - - @property - def logits(self): - return self._categorical.logits - - @property - def probs(self): - return self._categorical.probs - - def rsample(self, sample_shape=torch.Size()): + + def gumbel_softmax_sample(self, sample_shape=torch.Size()): shape = self._extended_shape(sample_shape) uniforms = clamp_probs( torch.rand(shape, dtype=self.logits.dtype, device=self.logits.device) ) gumbels = -((-(uniforms.log())).log()) - scores = (self.logits + gumbels) / self.temperature - return scores - scores.logsumexp(dim=-1, keepdim=True) - - def sample(self, sample_shape: torch.Size = torch.Size()) -> torch.Tensor: - """ - Generates a sample_shape shaped sample or sample_shape shaped batch of - samples if the distribution parameters are batched. - """ - with torch.no_grad(): - return self.rsample(sample_shape) - - def log_prob(self, value): - K = self._categorical._num_events - if self._validate_args: - self._validate_sample(value) - logits, value = broadcast_all(self.logits, value) - log_scale = torch.full_like( - self.temperature, float(K) - ).lgamma() - self.temperature.log().mul(-(K - 1)) - score = logits - value.mul(self.temperature) - score = (score - score.logsumexp(dim=-1, keepdim=True)).sum(-1) - return score + log_scale + y = self.logits + gumbels + return torch.nn.functional.softmax(y / self.temperature, dim=-1) - -class RelaxedOneHotCategorical(TransformedDistribution): - r""" - Creates a RelaxedOneHotCategorical distribution parametrized by - :attr:`temperature`, and either :attr:`probs` or :attr:`logits`. - This is a relaxed version of the :class:`OneHotCategorical` distribution, so - its samples are on simplex, and are reparametrizable. - - Example:: - - >>> # xdoctest: +IGNORE_WANT("non-deterministic") - >>> m = RelaxedOneHotCategorical(torch.tensor([2.2]), - ... torch.tensor([0.1, 0.2, 0.3, 0.4])) - >>> m.sample() - tensor([ 0.1294, 0.2324, 0.3859, 0.2523]) - - Args: - temperature (Tensor): relaxation temperature - probs (Tensor): event probabilities - logits (Tensor): unnormalized log probability for each event - """ - arg_constraints = {"probs": constraints.simplex, "logits": constraints.real_vector} - support = constraints.simplex - has_rsample = True - - def __init__(self, temperature, probs=None, logits=None, validate_args=None): - base_dist = ExpRelaxedCategorical( - temperature, probs, logits, validate_args=validate_args - ) - super().__init__(base_dist, ExpTransform(), validate_args=validate_args) - - def expand(self, batch_shape, _instance=None): - new = self._get_checked_instance(RelaxedOneHotCategorical, _instance) - return super().expand(batch_shape, _instance=new) - - @property - def temperature(self): - return self.base_dist.temperature - - @property - def logits(self): - return self.base_dist.logits - - @property - def probs(self): - return self.base_dist.probs + def rsample(self, sample_shape=torch.Size()): + samples = self.sample(sample_shape) + gumbel_softmax_samples = self.gumbel_softmax_sample(sample_shape) + return samples + (gumbel_softmax_samples - gumbel_softmax_samples.detach())