from typing import Union, List, Dict from collections import namedtuple import numpy as np import math import torch import torch.nn as nn import torch.nn.functional as F from ding.utils import list_split, MODEL_REGISTRY, squeeze, SequenceType from ding.torch_utils.network.diffusion import extract, cosine_beta_schedule, apply_conditioning, \ DiffusionUNet1d, TemporalValue Sample = namedtuple('Sample', 'trajectories values chains') def default_sample_fn(model, x, cond, t): b, *_, device = *x.shape, x.device model_mean, _, model_log_variance = model.p_mean_variance( x=x, cond=cond, t=t, ) noise = 0.5 * torch.randn_like(x) # no noise when t == 0 nonzero_mask = (1 - (t == 0).float()).reshape(b, *((1, ) * (len(x.shape) - 1))) values = torch.zeros(len(x), device=device) return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise, values def get_guide_output(guide, x, cond, t): x.requires_grad_() y = guide(x, cond, t).squeeze(dim=-1) grad = torch.autograd.grad([y.sum()], [x])[0] x.detach() return y, grad def n_step_guided_p_sample( model, x, cond, t, guide, scale=0.001, t_stopgrad=0, n_guide_steps=1, scale_grad_by_std=True, ): model_log_variance = extract(model.posterior_log_variance_clipped, t, x.shape) model_std = torch.exp(0.5 * model_log_variance) model_var = torch.exp(model_log_variance) for _ in range(n_guide_steps): with torch.enable_grad(): y, grad = get_guide_output(guide, x, cond, t) if scale_grad_by_std: grad = model_var * grad grad[t < t_stopgrad] = 0 x = x + scale * grad x = apply_conditioning(x, cond, model.action_dim) model_mean, _, model_log_variance = model.p_mean_variance(x=x, cond=cond, t=t) # no noise when t == 0 noise = torch.randn_like(x) noise[t == 0] = 0 return model_mean + model_std * noise, y class GaussianDiffusion(nn.Module): """ Overview: Gaussian diffusion model Arguments: - model (:obj:`str`): type of model - model_cfg (:obj:'dict') config of model - horizon (:obj:`int`): horizon of trajectory - obs_dim (:obj:`int`): Dim of the ovservation - action_dim (:obj:`int`): Dim of the ation - n_timesteps (:obj:`int`): Number of timesteps - predict_epsilon (:obj:'bool'): Whether predict epsilon - loss_discount (:obj:'float'): discount of loss - clip_denoised (:obj:'bool'): Whether use clip_denoised - action_weight (:obj:'float'): weight of action - loss_weights (:obj:'dict'): weight of loss """ def __init__( self, model: str, model_cfg: dict, horizon: int, obs_dim: Union[int, SequenceType], action_dim: Union[int, SequenceType], n_timesteps: int = 1000, predict_epsilon: bool = True, loss_discount: float = 1.0, clip_denoised: bool = False, action_weight: float = 1.0, loss_weights: dict = None, ) -> None: super().__init__() self.horizon = horizon self.obs_dim = obs_dim self.action_dim = action_dim self.transition_dim = obs_dim + action_dim if type(model) == str: model = eval(model) self.model = model(**model_cfg) self.predict_epsilon = predict_epsilon self.clip_denoised = clip_denoised betas = cosine_beta_schedule(n_timesteps) alphas = 1. - betas alphas_cumprod = torch.cumprod(alphas, axis=0) alphas_cumprod_prev = torch.cat([torch.ones(1), alphas_cumprod[:-1]]) self.n_timesteps = int(n_timesteps) self.register_buffer('betas', betas) self.register_buffer('alphas_cumprod', alphas_cumprod) self.register_buffer('alphas_cumprod_prev', alphas_cumprod_prev) # calculations for diffusion q(x_t | x_{t-1}) and others self.register_buffer('sqrt_alphas_cumprod', torch.sqrt(alphas_cumprod)) self.register_buffer('sqrt_one_minus_alphas_cumprod', torch.sqrt(1. - alphas_cumprod)) self.register_buffer('log_one_minus_alphas_cumprod', torch.log(1. - alphas_cumprod)) self.register_buffer('sqrt_recip_alphas_cumprod', torch.sqrt(1. / alphas_cumprod)) self.register_buffer('sqrt_recipm1_alphas_cumprod', torch.sqrt(1. / alphas_cumprod - 1)) # calculations for posterior q(x_{t-1} | x_t, x_0) posterior_variance = betas * (1. - alphas_cumprod_prev) / (1. - alphas_cumprod) self.register_buffer('posterior_variance', posterior_variance) # log calculation clipped because the posterior variance # is 0 at the beginning of the diffusion chain self.register_buffer('posterior_log_variance_clipped', torch.log(torch.clamp(posterior_variance, min=1e-20))) self.register_buffer('posterior_mean_coef1', betas * np.sqrt(alphas_cumprod_prev) / (1. - alphas_cumprod)) self.register_buffer( 'posterior_mean_coef2', (1. - alphas_cumprod_prev) * np.sqrt(alphas) / (1. - alphas_cumprod) ) self.loss_weights = self.get_loss_weights(action_weight, loss_discount, loss_weights) def get_loss_weights(self, action_weight: float, discount: float, weights_dict: dict): """ Overview: sets loss coefficients for trajectory Arguments: - action_weight (:obj:'float') coefficient on first action loss - discount (:obj:'float') multiplies t^th timestep of trajectory loss by discount**t - weights_dict (:obj:'dict') { i: c } multiplies dimension i of observation loss by c """ self.action_weight = action_weight dim_weights = torch.ones(self.transition_dim, dtype=torch.float32) # set loss coefficients for dimensions of observation if weights_dict is None: weights_dict = {} for ind, w in weights_dict.items(): dim_weights[self.action_dim + ind] *= w # decay loss with trajectory timestep: discount**t discounts = discount ** torch.arange(self.horizon, dtype=torch.float) discounts = discounts / discounts.mean() loss_weights = torch.einsum('h,t->ht', discounts, dim_weights) # manually set a0 weight loss_weights[0, :self.action_dim] = action_weight return loss_weights def predict_start_from_noise(self, x_t, t, noise): """ if self.predict_epsilon, model output is (scaled) noise; otherwise, model predicts x0 directly """ if self.predict_epsilon: return ( extract(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t - extract(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * noise ) else: return noise def q_posterior(self, x_start, x_t, t): """ Overview: give noise and step, compute mean, variance. Arguments: x_start (:obj:'tensor') noise trajectory in timestep 0 x_t (:obj:'tuple') noise trajectory in timestep t t (:obj:'int') timestep of diffusion step """ posterior_mean = ( extract(self.posterior_mean_coef1, t, x_t.shape) * x_start + extract(self.posterior_mean_coef2, t, x_t.shape) * x_t ) posterior_variance = extract(self.posterior_variance, t, x_t.shape) posterior_log_variance_clipped = extract(self.posterior_log_variance_clipped, t, x_t.shape) return posterior_mean, posterior_variance, posterior_log_variance_clipped def p_mean_variance(self, x, cond, t): x_recon = self.predict_start_from_noise(x, t=t, noise=self.model(x, cond, t)) if self.clip_denoised: x_recon.clamp_(-1., 1.) else: assert RuntimeError() model_mean, posterior_variance, posterior_log_variance = self.q_posterior(x_start=x_recon, x_t=x, t=t) return model_mean, posterior_variance, posterior_log_variance @torch.no_grad() def p_sample_loop(self, shape, cond, return_chain=False, sample_fn=default_sample_fn, plan_size=1, **sample_kwargs): device = self.betas.device batch_size = shape[0] x = torch.randn(shape, device=device) x = apply_conditioning(x, cond, self.action_dim) chain = [x] if return_chain else None for i in reversed(range(0, self.n_timesteps)): t = torch.full((batch_size, ), i, device=device, dtype=torch.long) x, values = sample_fn(self, x, cond, t, **sample_kwargs) x = apply_conditioning(x, cond, self.action_dim) if return_chain: chain.append(x) values = values.reshape(-1, plan_size, *values.shape[1:]) x = x.reshape(-1, plan_size, *x.shape[1:]) if plan_size > 1: inds = torch.argsort(values, dim=1, descending=True) x = x[torch.arange(x.size(0)).unsqueeze(1), inds] values = values[torch.arange(values.size(0)).unsqueeze(1), inds] if return_chain: chain = torch.stack(chain, dim=1) return Sample(x, values, chain) @torch.no_grad() def conditional_sample(self, cond, horizon=None, **sample_kwargs): """ conditions : [ (time, state), ... ] """ device = self.betas.device batch_size = len(cond[0]) horizon = horizon or self.horizon shape = (batch_size, horizon, self.transition_dim) return self.p_sample_loop(shape, cond, **sample_kwargs) def q_sample(self, x_start, t, noise=None): """ Arguments: conditions (:obj:'tuple') [ (time, state), ... ] conditions of diffusion t (:obj:'int') timestep of diffusion noise (:obj:'tensor.float') timestep's noise of diffusion """ if noise is None: noise = torch.randn_like(x_start) sample = ( extract(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start + extract(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise ) return sample def p_losses(self, x_start, cond, t): noise = torch.randn_like(x_start) x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise) x_noisy = apply_conditioning(x_noisy, cond, self.action_dim) x_recon = self.model(x_noisy, cond, t) x_recon = apply_conditioning(x_recon, cond, self.action_dim) assert noise.shape == x_recon.shape if self.predict_epsilon: loss = F.mse_loss(x_recon, noise, reduction='none') a0_loss = (loss[:, 0, :self.action_dim] / self.loss_weights[0, :self.action_dim].to(loss.device)).mean() loss = (loss * self.loss_weights.to(loss.device)).mean() else: loss = F.mse_loss(x_recon, x_start, reduction='none') a0_loss = (loss[:, 0, :self.action_dim] / self.loss_weights[0, :self.action_dim].to(loss.device)).mean() loss = (loss * self.loss_weights.to(loss.device)).mean() return loss, a0_loss def forward(self, cond, *args, **kwargs): return self.conditional_sample(cond, *args, **kwargs) class ValueDiffusion(GaussianDiffusion): """ Overview: Gaussian diffusion model for value function. """ def p_losses(self, x_start, cond, target, t): noise = torch.randn_like(x_start) x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise) x_noisy = apply_conditioning(x_noisy, cond, self.action_dim) pred = self.model(x_noisy, cond, t) loss = F.mse_loss(pred, target, reduction='none').mean() log = { 'mean_pred': pred.mean().item(), 'max_pred': pred.max().item(), 'min_pred': pred.min().item(), } return loss, log def forward(self, x, cond, t): return self.model(x, cond, t) @MODEL_REGISTRY.register('pd') class PlanDiffuser(nn.Module): """ Overview: Diffuser model for plan. Arguments: - diffuser_model (:obj:`str`): type of plan model - diffuser_model_cfg (:obj:'dict') config of diffuser_model - value_model (:obj:`str`): type of value model, if haven't use, set it as None - value_model_cfg (:obj:`int`): config of value_model - sample_kwargs : config of sample function """ def __init__( self, diffuser_model: str, diffuser_model_cfg: dict, value_model: str, value_model_cfg: dict, **sample_kwargs ): super().__init__() diffuser_model = eval(diffuser_model) self.diffuser = diffuser_model(**diffuser_model_cfg) self.value = None if value_model: value_model = eval(value_model) self.value = value_model(**value_model_cfg) self.sample_kwargs = sample_kwargs def diffuser_loss(self, x_start, cond, t): return self.diffuser.p_losses(x_start, cond, t) def value_loss(self, x_start, cond, target, t): return self.value.p_losses(x_start, cond, target, t) def get_eval(self, cond, batch_size=1): cond = self.repeat_cond(cond, batch_size) if self.value: samples = self.diffuser( cond, sample_fn=n_step_guided_p_sample, plan_size=batch_size, guide=self.value, **self.sample_kwargs ) # extract action [eval_num, batch_size, horizon, transition_dim] actions = samples.trajectories[:, :, :, :self.diffuser.action_dim] action = actions[:, 0, 0] return action else: samples = self.diffuser(cond, plan_size=batch_size) return samples.trajectories[:, :, :, self.diffuser.action_dim:].squeeze(1) def repeat_cond(self, cond, batch_size): for k, v in cond.items(): cond[k] = v.repeat_interleave(batch_size, dim=0) return cond @MODEL_REGISTRY.register('dd') class GaussianInvDynDiffusion(nn.Module): """ Overview: Gaussian diffusion model with Invdyn action model. Arguments: - model (:obj:`str`): type of model - model_cfg (:obj:'dict') config of model - horizon (:obj:`int`): horizon of trajectory - obs_dim (:obj:`int`): Dim of the ovservation - action_dim (:obj:`int`): Dim of the ation - n_timesteps (:obj:`int`): Number of timesteps - hidden_dim (:obj:'int'): hidden dim of inv_model - returns_condition (:obj:'bool'): Whether use returns condition - ar_inv (:obj:'bool'): Whether use inverse action learning - train_only_inv (:obj:'bool'): Whether train inverse action model only - predict_epsilon (:obj:'bool'): Whether predict epsilon - condition_guidance_w (:obj:'float'): weight of condition guidance - loss_discount (:obj:'float'): discount of loss """ def __init__( self, model: str, model_cfg: dict, horizon: int, obs_dim: Union[int, SequenceType], action_dim: Union[int, SequenceType], n_timesteps: int = 1000, hidden_dim: int = 256, returns_condition: bool = False, ar_inv: bool = False, train_only_inv: bool = False, predict_epsilon: bool = True, condition_guidance_w: float = 0.1, loss_discount: float = 1.0, clip_denoised: bool = False, ) -> None: super().__init__() self.horizon = horizon self.obs_dim = obs_dim self.action_dim = action_dim self.transition_dim = obs_dim + action_dim if type(model) == str: model = eval(model) self.model = model(**model_cfg) self.ar_inv = ar_inv self.train_only_inv = train_only_inv self.predict_epsilon = predict_epsilon self.condition_guidance_w = condition_guidance_w self.inv_model = nn.Sequential( nn.Linear(2 * self.obs_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, self.action_dim), ) self.returns_condition = returns_condition self.clip_denoised = clip_denoised betas = cosine_beta_schedule(n_timesteps) alphas = 1. - betas alphas_cumprod = torch.cumprod(alphas, axis=0) alphas_cumprod_prev = torch.cat([torch.ones(1), alphas_cumprod[:-1]]) self.n_timesteps = int(n_timesteps) self.register_buffer('betas', betas) self.register_buffer('alphas_cumprod', alphas_cumprod) self.register_buffer('alphas_cumprod_prev', alphas_cumprod_prev) # calculations for diffusion q(x_t | x_{t-1}) and others self.register_buffer('sqrt_alphas_cumprod', torch.sqrt(alphas_cumprod)) self.register_buffer('sqrt_one_minus_alphas_cumprod', torch.sqrt(1. - alphas_cumprod)) self.register_buffer('log_one_minus_alphas_cumprod', torch.log(1. - alphas_cumprod)) self.register_buffer('sqrt_recip_alphas_cumprod', torch.sqrt(1. / alphas_cumprod)) self.register_buffer('sqrt_recipm1_alphas_cumprod', torch.sqrt(1. / alphas_cumprod - 1)) # calculations for posterior q(x_{t-1} | x_t, x_0) posterior_variance = betas * (1. - alphas_cumprod_prev) / (1. - alphas_cumprod) self.register_buffer('posterior_variance', posterior_variance) # log calculation clipped because the posterior variance # is 0 at the beginning of the diffusion chain self.register_buffer('posterior_log_variance_clipped', torch.log(torch.clamp(posterior_variance, min=1e-20))) self.register_buffer('posterior_mean_coef1', betas * np.sqrt(alphas_cumprod_prev) / (1. - alphas_cumprod)) self.register_buffer( 'posterior_mean_coef2', (1. - alphas_cumprod_prev) * np.sqrt(alphas) / (1. - alphas_cumprod) ) self.loss_weights = self.get_loss_weights(loss_discount) def get_loss_weights(self, discount: int): self.action_weight = 1 dim_weights = torch.ones(self.obs_dim, dtype=torch.float32) # decay loss with trajectory timestep: discount**t discounts = discount ** torch.arange(self.horizon, dtype=torch.float) discounts = discounts / discounts.mean() loss_weights = torch.einsum('h,t->ht', discounts, dim_weights) # Cause things are conditioned on t=0 if self.predict_epsilon: loss_weights[0, :] = 0 return loss_weights def predict_start_from_noise(self, x_t, t, noise): """ if self.predict_epsilon, model output is (scaled) noise; otherwise, model predicts x0 directly """ if self.predict_epsilon: return ( extract(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t - extract(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * noise ) else: return noise def q_posterior(self, x_start, x_t, t): """ Arguments: x_start (:obj:'tensor') noise trajectory in timestep 0 x_t (:obj:'tuple') noise trajectory in timestep t t (:obj:'int') timestep of diffusion step """ posterior_mean = ( extract(self.posterior_mean_coef1, t, x_t.shape) * x_start + extract(self.posterior_mean_coef2, t, x_t.shape) * x_t ) posterior_variance = extract(self.posterior_variance, t, x_t.shape) posterior_log_variance_clipped = extract(self.posterior_log_variance_clipped, t, x_t.shape) return posterior_mean, posterior_variance, posterior_log_variance_clipped def p_mean_variance(self, x, cond, t, returns=None): """ Arguments: x (:obj:'tensor') noise trajectory in timestep t cond (:obj:'tuple') [ (time, state), ... ] state is init state of env, time = 0 t (:obj:'int') timestep of diffusion step returns (:obj:'tensor') condition returns of trajectory, returns is normal return returns: model_mean (:obj:'tensor.float') posterior_variance (:obj:'float') posterior_log_variance (:obj:'float') """ if self.returns_condition: # epsilon could be epsilon or x0 itself epsilon_cond = self.model(x, cond, t, returns, use_dropout=False) epsilon_uncond = self.model(x, cond, t, returns, force_dropout=True) epsilon = epsilon_uncond + self.condition_guidance_w * (epsilon_cond - epsilon_uncond) else: epsilon = self.model(x, cond, t) t = t.detach().to(torch.int64) x_recon = self.predict_start_from_noise(x, t=t, noise=epsilon) if self.clip_denoised: x_recon.clamp_(-1., 1.) else: assert RuntimeError() model_mean, posterior_variance, posterior_log_variance = self.q_posterior(x_start=x_recon, x_t=x, t=t) return model_mean, posterior_variance, posterior_log_variance @torch.no_grad() def p_sample(self, x, cond, t, returns=None): """ Arguments: x (:obj:'tensor') noise trajectory in timestep t cond (:obj:'tuple') [ (time, state), ... ] state is init state of env, time = 0 t (:obj:'int') timestep of diffusion step returns (:obj:'tensor') condition returns of trajectory, returns is normal return """ b, *_, device = *x.shape, x.device model_mean, _, model_log_variance = self.p_mean_variance(x=x, cond=cond, t=t, returns=returns) noise = 0.5 * torch.randn_like(x) # no noise when t == 0 nonzero_mask = (1 - (t == 0).float()).reshape(b, *((1, ) * (len(x.shape) - 1))) return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise @torch.no_grad() def p_sample_loop(self, shape, cond, returns=None, verbose=True, return_diffusion=False): """ Arguments: shape (:obj:'tuple') (batch_size, horizon, self.obs_dim) cond (:obj:'tuple') [ (time, state), ... ] state is init state of env, time = 0 returns (:obj:'tensor') condition returns of trajectory, returns is normal return horizon (:obj:'int') horizon of trajectory verbose (:obj:'bool') whether log diffusion progress return_diffusion (:obj:'bool') whether use return diffusion """ device = self.betas.device batch_size = shape[0] x = 0.5 * torch.randn(shape, device=device) # In this model, init state must be given by the env and without noise. x = apply_conditioning(x, cond, 0) if return_diffusion: diffusion = [x] for i in reversed(range(0, self.n_timesteps)): timesteps = torch.full((batch_size, ), i, device=device, dtype=torch.long) x = self.p_sample(x, cond, timesteps, returns) x = apply_conditioning(x, cond, 0) if return_diffusion: diffusion.append(x) if return_diffusion: return x, torch.stack(diffusion, dim=1) else: return x @torch.no_grad() def conditional_sample(self, cond, returns=None, horizon=None, *args, **kwargs): """ Arguments: conditions (:obj:'tuple') [ (time, state), ... ] state is init state of env, time is timestep of trajectory returns (:obj:'tensor') condition returns of trajectory, returns is normal return horizon (:obj:'int') horizon of trajectory returns: x (:obj:'tensor') tarjctory of env """ device = self.betas.device batch_size = len(cond[0]) horizon = horizon or self.horizon shape = (batch_size, horizon, self.obs_dim) return self.p_sample_loop(shape, cond, returns, *args, **kwargs) def q_sample(self, x_start, t, noise=None): """ Arguments: conditions (:obj:'tuple') [ (time, state), ... ] conditions of diffusion t (:obj:'int') timestep of diffusion noise (:obj:'tensor.float') timestep's noise of diffusion """ if noise is None: noise = torch.randn_like(x_start) sample = ( extract(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start + extract(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise ) return sample def p_losses(self, x_start, cond, t, returns=None): noise = torch.randn_like(x_start) x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise) x_noisy = apply_conditioning(x_noisy, cond, 0) x_recon = self.model(x_noisy, cond, t, returns) if not self.predict_epsilon: x_recon = apply_conditioning(x_recon, cond, 0) assert noise.shape == x_recon.shape if self.predict_epsilon: loss = F.mse_loss(x_recon, noise, reduction='none') loss = (loss * self.loss_weights.to(loss.device)).mean() else: loss = F.mse_loss(x_recon, x_start, reduction='none') loss = (loss * self.loss_weights.to(loss.device)).mean() return loss def forward(self, cond, *args, **kwargs): return self.conditional_sample(cond=cond, *args, **kwargs)