`ding.model.template.diffusion`¶

`ding.model.template.diffusion` ¶

`GaussianDiffusion` ¶

Bases: 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

`get_loss_weights(action_weight, discount, weights_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

`predict_start_from_noise(x_t, t, noise)` ¶

if self.predict_epsilon, model output is (scaled) noise; otherwise, model predicts x0 directly

`q_posterior(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

`conditional_sample(cond, horizon=None, **sample_kwargs)` ¶

conditions : [ (time, state), ... ]

`q_sample(x_start, t, noise=None)` ¶

Parameters:

Name	Description	Default
`conditions (`	obj:'tuple') [ (time, state), ... ] conditions of diffusion	required
`t (`	obj:'int') timestep of diffusion	required
`noise (`	obj:'tensor.float') timestep's noise of diffusion	required

`ValueDiffusion` ¶

Bases: GaussianDiffusion

Overview

Gaussian diffusion model for value function.

`PlanDiffuser` ¶

Bases: 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

`GaussianInvDynDiffusion` ¶

Bases: 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

`predict_start_from_noise(x_t, t, noise)` ¶

if self.predict_epsilon, model output is (scaled) noise; otherwise, model predicts x0 directly

`q_posterior(x_start, x_t, t)` ¶

Parameters:

Name	Description	Default
`x_start (`	obj:'tensor') noise trajectory in timestep 0	required
`x_t (`	obj:'tuple') noise trajectory in timestep t	required
`t (`	obj:'int') timestep of diffusion step	required

`p_mean_variance(x, cond, t, returns=None)` ¶

Parameters:

Name	Description	Default
`x (`	obj:'tensor') noise trajectory in timestep t	required
`cond (`	obj:'tuple') [ (time, state), ... ] state is init state of env, time = 0	required
`t (`	obj:'int') timestep of diffusion step	required
`returns (`	obj:'tensor') condition returns of trajectory, returns is normal return	required

returns: model_mean (:obj:'tensor.float') posterior_variance (:obj:'float') posterior_log_variance (:obj:'float')

`p_sample(x, cond, t, returns=None)` ¶

Parameters:

Name	Description	Default
`x (`	obj:'tensor') noise trajectory in timestep t	required
`cond (`	obj:'tuple') [ (time, state), ... ] state is init state of env, time = 0	required
`t (`	obj:'int') timestep of diffusion step	required
`returns (`	obj:'tensor') condition returns of trajectory, returns is normal return	required

`p_sample_loop(shape, cond, returns=None, verbose=True, return_diffusion=False)` ¶

Parameters:

Name	Description	Default
`shape (`	obj:'tuple') (batch_size, horizon, self.obs_dim)	required
`cond (`	obj:'tuple') [ (time, state), ... ] state is init state of env, time = 0	required
`returns (`	obj:'tensor') condition returns of trajectory, returns is normal return	required
`horizon (`	obj:'int') horizon of trajectory	required
`verbose (`	obj:'bool') whether log diffusion progress	required
`return_diffusion (`	obj:'bool') whether use return diffusion	required

`conditional_sample(cond, returns=None, horizon=None, *args, **kwargs)` ¶

Parameters:

Name	Description	Default
`conditions (`	obj:'tuple') [ (time, state), ... ] state is init state of env, time is timestep of trajectory	required
`returns (`	obj:'tensor') condition returns of trajectory, returns is normal return	required
`horizon (`	obj:'int') horizon of trajectory	required

returns: x (:obj:'tensor') tarjctory of env

`q_sample(x_start, t, noise=None)` ¶

Parameters:

Name	Description	Default
`conditions (`	obj:'tuple') [ (time, state), ... ] conditions of diffusion	required
`t (`	obj:'int') timestep of diffusion	required
`noise (`	obj:'tensor.float') timestep's noise of diffusion	required

Full Source Code

../ding/model/template/diffusion.py

from typing import Union, List, Dictfrom collections import namedtupleimport numpy as npimport mathimport torchimport torch.nn as nnimport torch.nn.functional as Ffrom ding.utils import list_split, MODEL_REGISTRY, squeeze, SequenceTypefrom ding.torch_utils.network.diffusion import extract, cosine_beta_schedule, apply_conditioning, \    DiffusionUNet1d, TemporalValueSample = 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, valuesdef 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, graddef 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, yclass 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)

ding.model.template.diffusion¶