DI-engine/ding/torch_utils/network/dreamer.py

import math
import numpy as np

import torch
from torch import nn
import torch.nn.functional as F
from torch import distributions as torchd
from ding.torch_utils import MLP
from ding.rl_utils import symlog, inv_symlog


class Conv2dSame(torch.nn.Conv2d):
    """
    Overview:
         Conv2dSame Network for dreamerv3.
    Interfaces:
        ``__init__``, ``forward``
    """

    def calc_same_pad(self, i, k, s, d):
        """
        Overview:
            Calculate the same padding size.
        Arguments:
            - i (:obj:`int`): Input size.
            - k (:obj:`int`): Kernel size.
            - s (:obj:`int`): Stride size.
            - d (:obj:`int`): Dilation size.
        """
        return max((math.ceil(i / s) - 1) * s + (k - 1) * d + 1 - i, 0)

    def forward(self, x):
        """
        Overview:
            compute the forward of Conv2dSame.
        Arguments:
            - x (:obj:`torch.Tensor`): Input tensor.
        """
        ih, iw = x.size()[-2:]
        pad_h = self.calc_same_pad(i=ih, k=self.kernel_size[0], s=self.stride[0], d=self.dilation[0])
        pad_w = self.calc_same_pad(i=iw, k=self.kernel_size[1], s=self.stride[1], d=self.dilation[1])

        if pad_h > 0 or pad_w > 0:
            x = F.pad(x, [pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2])

        ret = F.conv2d(
            x,
            self.weight,
            self.bias,
            self.stride,
            self.padding,
            self.dilation,
            self.groups,
        )
        return ret


class DreamerLayerNorm(nn.Module):
    """
    Overview:
         DreamerLayerNorm Network for dreamerv3.
    Interfaces:
        ``__init__``, ``forward``
    """

    def __init__(self, ch, eps=1e-03):
        """
        Overview:
            Init the DreamerLayerNorm class.
        Arguments:
            - ch (:obj:`int`): Input channel.
            - eps (:obj:`float`): Epsilon.
        """

        super(DreamerLayerNorm, self).__init__()
        self.norm = torch.nn.LayerNorm(ch, eps=eps)

    def forward(self, x):
        """
        Overview:
            compute the forward of DreamerLayerNorm.
        Arguments:
            - x (:obj:`torch.Tensor`): Input tensor.
        """

        x = x.permute(0, 2, 3, 1)
        x = self.norm(x)
        x = x.permute(0, 3, 1, 2)
        return x


class DenseHead(nn.Module):
    """
    Overview:
       DenseHead Network for value head, reward head, and discount head of dreamerv3.
    Interfaces:
        ``__init__``, ``forward``
    """

    def __init__(
        self,
        inp_dim,
        shape,  # (255,)
        layer_num,
        units,  # 512
        act='SiLU',
        norm='LN',
        dist='normal',
        std=1.0,
        outscale=1.0,
        device='cpu',
    ):
        """
        Overview:
            Init the DenseHead class.
        Arguments:
            - inp_dim (:obj:`int`): Input dimension.
            - shape (:obj:`tuple`): Output shape.
            - layer_num (:obj:`int`): Number of layers.
            - units (:obj:`int`): Number of units.
            - act (:obj:`str`): Activation function.
            - norm (:obj:`str`): Normalization function.
            - dist (:obj:`str`): Distribution function.
            - std (:obj:`float`): Standard deviation.
            - outscale (:obj:`float`): Output scale.
            - device (:obj:`str`): Device.
        """

        super(DenseHead, self).__init__()
        self._shape = (shape, ) if isinstance(shape, int) else shape
        if len(self._shape) == 0:
            self._shape = (1, )
        self._layer_num = layer_num
        self._units = units
        self._act = getattr(torch.nn, act)()
        self._norm = norm
        self._dist = dist
        self._std = std
        self._device = device

        self.mlp = MLP(
            inp_dim,
            self._units,
            self._units,
            self._layer_num,
            layer_fn=nn.Linear,
            activation=self._act,
            norm_type=self._norm
        )
        self.mlp.apply(weight_init)

        self.mean_layer = nn.Linear(self._units, np.prod(self._shape))
        self.mean_layer.apply(uniform_weight_init(outscale))

        if self._std == "learned":
            self.std_layer = nn.Linear(self._units, np.prod(self._shape))
            self.std_layer.apply(uniform_weight_init(outscale))

    def forward(self, features):
        """
        Overview:
            compute the forward of DenseHead.
        Arguments:
            - features (:obj:`torch.Tensor`): Input tensor.
        """

        x = features
        out = self.mlp(x)  # (batch, time, _units=512)
        mean = self.mean_layer(out)  # (batch, time, 255)
        if self._std == "learned":
            std = self.std_layer(out)
        else:
            std = self._std
        if self._dist == "normal":
            return ContDist(torchd.independent.Independent(torchd.normal.Normal(mean, std), len(self._shape)))
        elif self._dist == "huber":
            return ContDist(torchd.independent.Independent(UnnormalizedHuber(mean, std, 1.0), len(self._shape)))
        elif self._dist == "binary":
            return Bernoulli(torchd.independent.Independent(torchd.bernoulli.Bernoulli(logits=mean), len(self._shape)))
        elif self._dist == "twohot_symlog":
            return TwoHotDistSymlog(logits=mean, device=self._device)
        raise NotImplementedError(self._dist)


class ActionHead(nn.Module):
    """
    Overview:
       ActionHead Network for action head of dreamerv3.
    Interfaces:
        ``__init__``, ``forward``
    """

    def __init__(
        self,
        inp_dim,
        size,
        layers,
        units,
        act=nn.ELU,
        norm=nn.LayerNorm,
        dist="trunc_normal",
        init_std=0.0,
        min_std=0.1,
        max_std=1.0,
        temp=0.1,
        outscale=1.0,
        unimix_ratio=0.01,
    ):
        """
        Overview:
            Initialize the ActionHead class.
        Arguments:
            - inp_dim (:obj:`int`): Input dimension.
            - size (:obj:`int`): Output size.
            - layers (:obj:`int`): Number of layers.
            - units (:obj:`int`): Number of units.
            - act (:obj:`str`): Activation function.
            - norm (:obj:`str`): Normalization function.
            - dist (:obj:`str`): Distribution function.
            - init_std (:obj:`float`): Initial standard deviation.
            - min_std (:obj:`float`): Minimum standard deviation.
            - max_std (:obj:`float`): Maximum standard deviation.
            - temp (:obj:`float`): Temperature.
            - outscale (:obj:`float`): Output scale.
            - unimix_ratio (:obj:`float`): Unimix ratio.
        """
        super(ActionHead, self).__init__()
        self._size = size
        self._layers = layers
        self._units = units
        self._dist = dist
        self._act = getattr(torch.nn, act)
        self._norm = getattr(torch.nn, norm)
        self._min_std = min_std
        self._max_std = max_std
        self._init_std = init_std
        self._unimix_ratio = unimix_ratio
        self._temp = temp() if callable(temp) else temp

        pre_layers = []
        for index in range(self._layers):
            pre_layers.append(nn.Linear(inp_dim, self._units, bias=False))
            pre_layers.append(self._norm(self._units, eps=1e-03))
            pre_layers.append(self._act())
            if index == 0:
                inp_dim = self._units
        self._pre_layers = nn.Sequential(*pre_layers)
        self._pre_layers.apply(weight_init)

        if self._dist in ["tanh_normal", "tanh_normal_5", "normal", "trunc_normal"]:
            self._dist_layer = nn.Linear(self._units, 2 * self._size)
            self._dist_layer.apply(uniform_weight_init(outscale))

        elif self._dist in ["normal_1", "onehot", "onehot_gumbel"]:
            self._dist_layer = nn.Linear(self._units, self._size)
            self._dist_layer.apply(uniform_weight_init(outscale))

    def forward(self, features):
        """
        Overview:
            compute the forward of ActionHead.
        Arguments:
            - features (:obj:`torch.Tensor`): Input tensor.
        """

        x = features
        x = self._pre_layers(x)
        if self._dist == "tanh_normal":
            x = self._dist_layer(x)
            mean, std = torch.split(x, 2, -1)
            mean = torch.tanh(mean)
            std = F.softplus(std + self._init_std) + self._min_std
            dist = torchd.normal.Normal(mean, std)
            dist = torchd.transformed_distribution.TransformedDistribution(dist, TanhBijector())
            dist = torchd.independent.Independent(dist, 1)
            dist = SampleDist(dist)
        elif self._dist == "tanh_normal_5":
            x = self._dist_layer(x)
            mean, std = torch.split(x, 2, -1)
            mean = 5 * torch.tanh(mean / 5)
            std = F.softplus(std + 5) + 5
            dist = torchd.normal.Normal(mean, std)
            dist = torchd.transformed_distribution.TransformedDistribution(dist, TanhBijector())
            dist = torchd.independent.Independent(dist, 1)
            dist = SampleDist(dist)
        elif self._dist == "normal":
            x = self._dist_layer(x)
            mean, std = torch.split(x, [self._size] * 2, -1)
            std = (self._max_std - self._min_std) * torch.sigmoid(std + 2.0) + self._min_std
            dist = torchd.normal.Normal(torch.tanh(mean), std)
            dist = ContDist(torchd.independent.Independent(dist, 1))
        elif self._dist == "normal_1":
            x = self._dist_layer(x)
            dist = torchd.normal.Normal(mean, 1)
            dist = ContDist(torchd.independent.Independent(dist, 1))
        elif self._dist == "trunc_normal":
            x = self._dist_layer(x)
            mean, std = torch.split(x, [self._size] * 2, -1)
            mean = torch.tanh(mean)
            std = 2 * torch.sigmoid(std / 2) + self._min_std
            dist = SafeTruncatedNormal(mean, std, -1, 1)
            dist = ContDist(torchd.independent.Independent(dist, 1))
        elif self._dist == "onehot":
            x = self._dist_layer(x)
            dist = OneHotDist(x, unimix_ratio=self._unimix_ratio)
        elif self._dist == "onehot_gumble":
            x = self._dist_layer(x)
            temp = self._temp
            dist = ContDist(torchd.gumbel.Gumbel(x, 1 / temp))
        else:
            raise NotImplementedError(self._dist)
        return dist


class SampleDist:
    """
    Overview:
       A kind of sample Dist for ActionHead of dreamerv3.
    Interfaces:
        ``__init__``, ``mean``, ``mode``, ``entropy``
    """

    def __init__(self, dist, samples=100):
        """
        Overview:
            Initialize the SampleDist class.
        Arguments:
            - dist (:obj:`torch.Tensor`): Distribution.
            - samples (:obj:`int`): Number of samples.
        """

        self._dist = dist
        self._samples = samples

    def mean(self):
        """
        Overview:
            Calculate the mean of the distribution.
        """

        samples = self._dist.sample(self._samples)
        return torch.mean(samples, 0)

    def mode(self):
        """
        Overview:
            Calculate the mode of the distribution.
        """

        sample = self._dist.sample(self._samples)
        logprob = self._dist.log_prob(sample)
        return sample[torch.argmax(logprob)][0]

    def entropy(self):
        """
        Overview:
            Calculate the entropy of the distribution.
        """

        sample = self._dist.sample(self._samples)
        logprob = self.log_prob(sample)
        return -torch.mean(logprob, 0)


class OneHotDist(torchd.one_hot_categorical.OneHotCategorical):
    """
    Overview:
       A kind of onehot Dist for dreamerv3.
    Interfaces:
        ``__init__``, ``mode``, ``sample``
    """

    def __init__(self, logits=None, probs=None, unimix_ratio=0.0):
        """
        Overview:
            Initialize the OneHotDist class.
        Arguments:
            - logits (:obj:`torch.Tensor`): Logits.
            - probs (:obj:`torch.Tensor`): Probabilities.
            - unimix_ratio (:obj:`float`): Unimix ratio.
        """

        if logits is not None and unimix_ratio > 0.0:
            probs = F.softmax(logits, dim=-1)
            probs = probs * (1.0 - unimix_ratio) + unimix_ratio / probs.shape[-1]
            logits = torch.log(probs)
            super().__init__(logits=logits, probs=None)
        else:
            super().__init__(logits=logits, probs=probs)

    def mode(self):
        """
        Overview:
            Calculate the mode of the distribution.
        """

        _mode = F.one_hot(torch.argmax(super().logits, axis=-1), super().logits.shape[-1])
        return _mode.detach() + super().logits - super().logits.detach()

    def sample(self, sample_shape=(), seed=None):
        """
        Overview:
            Sample from the distribution.
        Arguments:
            - sample_shape (:obj:`tuple`): Sample shape.
            - seed (:obj:`int`): Seed.
        """

        if seed is not None:
            raise ValueError('need to check')
        sample = super().sample(sample_shape)
        probs = super().probs
        while len(probs.shape) < len(sample.shape):
            probs = probs[None]
        sample += probs - probs.detach()
        return sample


class TwoHotDistSymlog:
    """
    Overview:
       A kind of twohotsymlog Dist for dreamerv3.
    Interfaces:
        ``__init__``, ``mode``, ``mean``, ``log_prob``, ``log_prob_target``
    """

    def __init__(self, logits=None, low=-20.0, high=20.0, device='cpu'):
        """
        Overview:
            Initialize the TwoHotDistSymlog class.
        Arguments:
            - logits (:obj:`torch.Tensor`): Logits.
            - low (:obj:`float`): Low.
            - high (:obj:`float`): High.
            - device (:obj:`str`): Device.
        """

        self.logits = logits
        self.probs = torch.softmax(logits, -1)
        self.buckets = torch.linspace(low, high, steps=255).to(device)
        self.width = (self.buckets[-1] - self.buckets[0]) / 255

    def mean(self):
        """
        Overview:
            Calculate the mean of the distribution.
        """

        _mean = self.probs * self.buckets
        return inv_symlog(torch.sum(_mean, dim=-1, keepdim=True))

    def mode(self):
        """
        Overview:
            Calculate the mode of the distribution.
        """

        _mode = self.probs * self.buckets
        return inv_symlog(torch.sum(_mode, dim=-1, keepdim=True))

    # Inside OneHotCategorical, log_prob is calculated using only max element in targets
    def log_prob(self, x):
        """
        Overview:
            Calculate the log probability of the distribution.
        Arguments:
            - x (:obj:`torch.Tensor`): Input tensor.
        """

        x = symlog(x)
        # x(time, batch, 1)
        below = torch.sum((self.buckets <= x[..., None]).to(torch.int32), dim=-1) - 1
        above = len(self.buckets) - torch.sum((self.buckets > x[..., None]).to(torch.int32), dim=-1)
        below = torch.clip(below, 0, len(self.buckets) - 1)
        above = torch.clip(above, 0, len(self.buckets) - 1)
        equal = (below == above)

        dist_to_below = torch.where(equal, 1, torch.abs(self.buckets[below] - x))
        dist_to_above = torch.where(equal, 1, torch.abs(self.buckets[above] - x))
        total = dist_to_below + dist_to_above
        weight_below = dist_to_above / total
        weight_above = dist_to_below / total
        target = (
            F.one_hot(below, num_classes=len(self.buckets)) * weight_below[..., None] +
            F.one_hot(above, num_classes=len(self.buckets)) * weight_above[..., None]
        )
        log_pred = self.logits - torch.logsumexp(self.logits, -1, keepdim=True)
        target = target.squeeze(-2)

        return (target * log_pred).sum(-1)

    def log_prob_target(self, target):
        """
        Overview:
            Calculate the log probability of the target.
        Arguments:
            - target (:obj:`torch.Tensor`): Target tensor.
        """

        log_pred = super().logits - torch.logsumexp(super().logits, -1, keepdim=True)
        return (target * log_pred).sum(-1)


class SymlogDist:
    """
    Overview:
       A kind of Symlog Dist for dreamerv3.
    Interfaces:
        ``__init__``, ``entropy``, ``mode``, ``mean``, ``log_prob``
    """

    def __init__(self, mode, dist='mse', aggregation='sum', tol=1e-8, dim_to_reduce=[-1, -2, -3]):
        """
        Overview:
            Initialize the SymlogDist class.
        Arguments:
            - mode (:obj:`torch.Tensor`): Mode.
            - dist (:obj:`str`): Distribution function.
            - aggregation (:obj:`str`): Aggregation function.
            - tol (:obj:`float`): Tolerance.
            - dim_to_reduce (:obj:`list`): Dimension to reduce.
        """
        self._mode = mode
        self._dist = dist
        self._aggregation = aggregation
        self._tol = tol
        self._dim_to_reduce = dim_to_reduce

    def mode(self):
        """
        Overview:
            Calculate the mode of the distribution.
        """

        return inv_symlog(self._mode)

    def mean(self):
        """
        Overview:
            Calculate the mean of the distribution.
        """

        return inv_symlog(self._mode)

    def log_prob(self, value):
        """
        Overview:
            Calculate the log probability of the distribution.
        Arguments:
            - value (:obj:`torch.Tensor`): Input tensor.
        """

        assert self._mode.shape == value.shape
        if self._dist == 'mse':
            distance = (self._mode - symlog(value)) ** 2.0
            distance = torch.where(distance < self._tol, 0, distance)
        elif self._dist == 'abs':
            distance = torch.abs(self._mode - symlog(value))
            distance = torch.where(distance < self._tol, 0, distance)
        else:
            raise NotImplementedError(self._dist)
        if self._aggregation == 'mean':
            loss = distance.mean(self._dim_to_reduce)
        elif self._aggregation == 'sum':
            loss = distance.sum(self._dim_to_reduce)
        else:
            raise NotImplementedError(self._aggregation)
        return -loss


class ContDist:
    """
    Overview:
       A kind of ordinary Dist for dreamerv3.
    Interfaces:
        ``__init__``, ``entropy``, ``mode``, ``sample``, ``log_prob``
    """

    def __init__(self, dist=None):
        """
        Overview:
            Initialize the ContDist class.
        Arguments:
            - dist (:obj:`torch.Tensor`): Distribution.
        """

        super().__init__()
        self._dist = dist
        self.mean = dist.mean

    def __getattr__(self, name):
        """
        Overview:
            Get attribute.
        Arguments:
            - name (:obj:`str`): Attribute name.
        """

        return getattr(self._dist, name)

    def entropy(self):
        """
        Overview:
            Calculate the entropy of the distribution.
        """

        return self._dist.entropy()

    def mode(self):
        """
        Overview:
            Calculate the mode of the distribution.
        """

        return self._dist.mean

    def sample(self, sample_shape=()):
        """
        Overview:
            Sample from the distribution.
        Arguments:
            - sample_shape (:obj:`tuple`): Sample shape.
        """

        return self._dist.rsample(sample_shape)

    def log_prob(self, x):
        return self._dist.log_prob(x)


class Bernoulli:
    """
    Overview:
       A kind of Bernoulli Dist for dreamerv3.
    Interfaces:
        ``__init__``, ``entropy``, ``mode``, ``sample``, ``log_prob``
    """

    def __init__(self, dist=None):
        """
        Overview:
            Initialize the Bernoulli distribution.
        Arguments:
            - dist (:obj:`torch.Tensor`): Distribution.
        """

        super().__init__()
        self._dist = dist
        self.mean = dist.mean

    def __getattr__(self, name):
        """
        Overview:
            Get attribute.
        Arguments:
            - name (:obj:`str`): Attribute name.
        """

        return getattr(self._dist, name)

    def entropy(self):
        """
        Overview:
            Calculate the entropy of the distribution.
        """
        return self._dist.entropy()

    def mode(self):
        """
        Overview:
            Calculate the mode of the distribution.
        """

        _mode = torch.round(self._dist.mean)
        return _mode.detach() + self._dist.mean - self._dist.mean.detach()

    def sample(self, sample_shape=()):
        """
        Overview:
            Sample from the distribution.
        Arguments:
            - sample_shape (:obj:`tuple`): Sample shape.
        """

        return self._dist.rsample(sample_shape)

    def log_prob(self, x):
        """
        Overview:
            Calculate the log probability of the distribution.
        Arguments:
            - x (:obj:`torch.Tensor`): Input tensor.
        """

        _logits = self._dist.base_dist.logits
        log_probs0 = -F.softplus(_logits)
        log_probs1 = -F.softplus(-_logits)

        return log_probs0 * (1 - x) + log_probs1 * x


class UnnormalizedHuber(torchd.normal.Normal):
    """
    Overview:
       A kind of UnnormalizedHuber Dist for dreamerv3.
    Interfaces:
        ``__init__``, ``mode``, ``log_prob``
    """

    def __init__(self, loc, scale, threshold=1, **kwargs):
        """
        Overview:
            Initialize the UnnormalizedHuber class.
        Arguments:
            - loc (:obj:`torch.Tensor`): Location.
            - scale (:obj:`torch.Tensor`): Scale.
            - threshold (:obj:`float`): Threshold.
        """
        super().__init__(loc, scale, **kwargs)
        self._threshold = threshold

    def log_prob(self, event):
        """
        Overview:
            Calculate the log probability of the distribution.
        Arguments:
            - event (:obj:`torch.Tensor`): Event.
        """

        return -(torch.sqrt((event - self.mean) ** 2 + self._threshold ** 2) - self._threshold)

    def mode(self):
        """
        Overview:
            Calculate the mode of the distribution.
        """

        return self.mean


class SafeTruncatedNormal(torchd.normal.Normal):
    """
    Overview:
       A kind of SafeTruncatedNormal Dist for dreamerv3.
    Interfaces:
        ``__init__``, ``sample``
    """

    def __init__(self, loc, scale, low, high, clip=1e-6, mult=1):
        """
        Overview:
            Initialize the SafeTruncatedNormal class.
        Arguments:
            - loc (:obj:`torch.Tensor`): Location.
            - scale (:obj:`torch.Tensor`): Scale.
            - low (:obj:`float`): Low.
            - high (:obj:`float`): High.
            - clip (:obj:`float`): Clip.
            - mult (:obj:`float`): Mult.
        """

        super().__init__(loc, scale)
        self._low = low
        self._high = high
        self._clip = clip
        self._mult = mult

    def sample(self, sample_shape):
        """
        Overview:
            Sample from the distribution.
        Arguments:
            - sample_shape (:obj:`tuple`): Sample shape.
        """

        event = super().sample(sample_shape)
        if self._clip:
            clipped = torch.clip(event, self._low + self._clip, self._high - self._clip)
            event = event - event.detach() + clipped.detach()
        if self._mult:
            event *= self._mult
        return event


class TanhBijector(torchd.Transform):
    """
    Overview:
       A kind of TanhBijector Dist for dreamerv3.
    Interfaces:
        ``__init__``, ``_forward``, ``_inverse``, ``_forward_log_det_jacobian``
    """

    def __init__(self, validate_args=False, name='tanh'):
        """
        Overview:
            Initialize the TanhBijector class.
        Arguments:
            - validate_args (:obj:`bool`): Validate arguments.
            - name (:obj:`str`): Name.
        """

        super().__init__()

    def _forward(self, x):
        """
        Overview:
            Calculate the forward of the distribution.
        Arguments:
            - x (:obj:`torch.Tensor`): Input tensor.
        """

        return torch.tanh(x)

    def _inverse(self, y):
        """
        Overview:
            Calculate the inverse of the distribution.
        Arguments:
            - y (:obj:`torch.Tensor`): Input tensor.
        """

        y = torch.where((torch.abs(y) <= 1.), torch.clamp(y, -0.99999997, 0.99999997), y)
        y = torch.atanh(y)
        return y

    def _forward_log_det_jacobian(self, x):
        """
        Overview:
            Calculate the forward log det jacobian of the distribution.
        Arguments:
            - x (:obj:`torch.Tensor`): Input tensor.
        """

        log2 = torch.math.log(2.0)
        return 2.0 * (log2 - x - torch.softplus(-2.0 * x))


def static_scan(fn, inputs, start):
    """
    Overview:
         Static scan function.
    Arguments:
        - fn (:obj:`function`): Function.
        - inputs (:obj:`tuple`): Inputs.
        - start (:obj:`torch.Tensor`): Start tensor.
    """

    last = start  # {logit, stoch, deter:[batch_size, self._deter]}
    indices = range(inputs[0].shape[0])
    flag = True
    for index in indices:
        inp = lambda x: (_input[x] for _input in inputs)  # inputs:(action:(time, batch, 6), embed:(time, batch, 4096))
        last = fn(last, *inp(index))  # post, prior
        if flag:
            if isinstance(last, dict):
                outputs = {key: value.clone().unsqueeze(0) for key, value in last.items()}
            else:
                outputs = []
                for _last in last:
                    if isinstance(_last, dict):
                        outputs.append({key: value.clone().unsqueeze(0) for key, value in _last.items()})
                    else:
                        outputs.append(_last.clone().unsqueeze(0))
            flag = False
        else:
            if isinstance(last, dict):
                for key in last.keys():
                    outputs[key] = torch.cat([outputs[key], last[key].unsqueeze(0)], dim=0)
            else:
                for j in range(len(outputs)):
                    if isinstance(last[j], dict):
                        for key in last[j].keys():
                            outputs[j][key] = torch.cat([outputs[j][key], last[j][key].unsqueeze(0)], dim=0)
                    else:
                        outputs[j] = torch.cat([outputs[j], last[j].unsqueeze(0)], dim=0)
    if isinstance(last, dict):
        outputs = [outputs]
    return outputs


def weight_init(m):
    """
    Overview:
       weight_init for Linear, Conv2d, ConvTranspose2d, and LayerNorm.
    Arguments:
        - m (:obj:`torch.nn`): Module.
    """

    if isinstance(m, nn.Linear):
        in_num = m.in_features
        out_num = m.out_features
        denoms = (in_num + out_num) / 2.0
        scale = 1.0 / denoms
        std = np.sqrt(scale) / 0.87962566103423978
        nn.init.trunc_normal_(m.weight.data, mean=0.0, std=std, a=-2.0, b=2.0)
        if hasattr(m.bias, 'data'):
            m.bias.data.fill_(0.0)
    elif isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d):
        space = m.kernel_size[0] * m.kernel_size[1]
        in_num = space * m.in_channels
        out_num = space * m.out_channels
        denoms = (in_num + out_num) / 2.0
        scale = 1.0 / denoms
        std = np.sqrt(scale) / 0.87962566103423978
        nn.init.trunc_normal_(m.weight.data, mean=0.0, std=std, a=-2.0, b=2.0)
        if hasattr(m.bias, 'data'):
            m.bias.data.fill_(0.0)
    elif isinstance(m, nn.LayerNorm):
        m.weight.data.fill_(1.0)
        if hasattr(m.bias, 'data'):
            m.bias.data.fill_(0.0)


def uniform_weight_init(given_scale):
    """
    Overview:
       weight_init for Linear and LayerNorm.
    Arguments:
        - given_scale (:obj:`float`): Given scale.
    """

    def f(m):
        if isinstance(m, nn.Linear):
            in_num = m.in_features
            out_num = m.out_features
            denoms = (in_num + out_num) / 2.0
            scale = given_scale / denoms
            limit = np.sqrt(3 * scale)
            nn.init.uniform_(m.weight.data, a=-limit, b=limit)
            if hasattr(m.bias, 'data'):
                m.bias.data.fill_(0.0)
        elif isinstance(m, nn.LayerNorm):
            m.weight.data.fill_(1.0)
            if hasattr(m.bias, 'data'):
                m.bias.data.fill_(0.0)

    return f