init test project

GPS.py

@@ -0,0 +1,324 @@
+import os
+import os.path as osp
+import torch
+import torch.nn.functional as F
+import numpy as np
+import itertools
+from tensorboardX import SummaryWriter
+
+from NN.losses import make_criteria
+from utils.base import logger
+
+class GPS:
+    def __init__(self,
+                 init_mode: str = 'random_synthesis',
+                 noise_sigma: float = 1.0,
+                 coarse_ratio: float = 0.2,
+                 coarse_ratio_factor: float = 6,
+                 pyr_factor: float = 0.75,
+                 num_stages_limit: int = -1,
+                 device: str = 'cuda:0',
+                 silent: bool = False
+                 ):
+        '''
+        Args:
+            init_mode:
+                - 'random_synthesis': init with random seed
+                - 'random': init with random seed
+            noise_sigma: float = 1.0, random noise.
+            coarse_ratio: float = 0.2, ratio at the coarse level.
+            pyr_factor: float = 0.75, pyramid factor.
+            num_stages_limit: int = -1, no limit.
+            device: str = 'cuda:0', default device.
+            silent: bool = False, mute the output.
+        '''
+        self.init_mode = init_mode
+        self.noise_sigma = noise_sigma
+        self.coarse_ratio = coarse_ratio
+        self.coarse_ratio_factor = coarse_ratio_factor
+        self.pyr_factor = pyr_factor
+        self.num_stages_limit = num_stages_limit
+        self.device = torch.device(device)
+        self.silent = silent
+
+    def _get_pyramid_lengths(self, dest, ext=None):
+        """Get a list of pyramid lengths"""
+        if self.coarse_ratio == -1:
+            self.coarse_ratio = np.around(ext['criteria']['patch_size'] * self.coarse_ratio_factor / dest, 2)
+
+        lengths = [int(np.round(dest * self.coarse_ratio))]
+        while lengths[-1] < dest:
+            lengths.append(int(np.round(lengths[-1] / self.pyr_factor)))
+            if lengths[-1] == lengths[-2]:
+                lengths[-1] += 1
+        lengths[-1] = dest
+
+        return lengths
+
+    def _get_target_pyramid(self, target, ext=None):
+        """Reads a target motion(s) and create a pyraimd out of it. Ordered in increatorch.sing size"""
+        self._num_target = len(target)
+        lengths = []
+        min_len = 10000
+        for i in range(len(target)):
+            new_length = self._get_pyramid_lengths(len(target[i]), ext)
+            min_len = min(min_len, len(new_length))
+            if self.num_stages_limit != -1:
+                new_length = new_length[:self.num_stages_limit]
+            lengths.append(new_length)
+        for i in range(len(target)):
+            lengths[i] = lengths[i][-min_len:]
+        self.pyraimd_lengths = lengths
+
+        target_pyramid = [[] for _ in range(len(lengths[0]))]
+        for step in range(len(lengths[0])):
+            for i in range(len(target)):
+                length = lengths[i][step]
+                motion = target[i]
+                target_pyramid[step].append(motion.sample(size=length).to(self.device))
+                # target_pyramid[step].append(motion.pos2velo(motion.sample(size=length)))
+                # motion.motion_data = motion.pos2velo(motion.motion_data)
+                # target_pyramid[step].append(motion.sample(size=length))
+                # motion.motion_data = motion.velo2pos(motion.motion_data)    
+
+        if not self.silent:
+            print('Levels:', lengths)
+            for i in range(len(target_pyramid)):
+                print(f'Number of clips in target pyramid {i} is {len(target_pyramid[i])}: {[[tgt.min(), tgt.max()] for tgt in target_pyramid[i]]}')
+
+        return target_pyramid
+
+    def _get_initial_motion(self):
+        """Prepare the initial motion for optimization"""
+        if 'random_synthesis' in str(self.init_mode):
+            m = self.init_mode.split('/')[-1]
+            if m =='random_synthesis':
+                final_length = sum([i[-1] for i in self.pyraimd_lengths])
+            elif 'x' in m:
+                final_length = int(m.replace('x', '')) * sum([i[-1] for i in self.pyraimd_lengths])
+            elif (self.init_mode.split('/')[-1]).isdigit():
+                final_length = int(self.init_mode.split('/')[-1])
+            else:
+                raise ValueError(f'incorrect init_mode: {self.init_mode}')
+
+            self.synthesized_lengths = self._get_pyramid_lengths(final_length)
+
+        else:
+            raise ValueError(f'Unsupported init_mode {self.init_mode}')
+            
+        initial_motion = F.interpolate(torch.cat([self.target_pyramid[0][i] for i in range(self._num_target)], dim=-1),
+                                       size=self.synthesized_lengths[0], mode='linear', align_corners=True)
+        if self.noise_sigma > 0:
+            initial_motion_w_noise = initial_motion + torch.randn_like(initial_motion) * self.noise_sigma
+            initial_motion_w_noise = torch.fmod(initial_motion_w_noise, 1.0)
+        else:
+            initial_motion_w_noise = initial_motion
+
+        if not self.silent:
+            print('Synthesized lengths:', self.synthesized_lengths)
+            print('Initial motion:', initial_motion.min(), initial_motion.max())
+            print('Initial motion with noise:', initial_motion_w_noise.min(), initial_motion_w_noise.max())
+
+        return initial_motion_w_noise
+
+    def run(self, target, mode="backpropagate", ext=None, debug_dir=None):
+        '''
+        Run the patch-based motion synthesis.
+
+        Args:
+            target (torch.Tensor): Target data.
+            mode (str): Optimization mode. Support ['backpropagate', 'match_and_blend']
+            ext (dict): extra data or constrain.
+            debug_dir (str): Debug directory.
+        '''
+        # preprare data
+        self.target_pyramid = self._get_target_pyramid(target, ext)
+        self.synthesized = self._get_initial_motion()
+        if debug_dir is not None:
+            writer = SummaryWriter(log_dir=debug_dir)
+
+        # prepare configuration
+        if mode == "backpropagate":
+            self.synthesized.requires_grad_(True)
+            assert 'criteria' in ext.keys(), 'Please specify a criteria for synthsis.'
+            criteria = make_criteria(ext['criteria']).to(self.device)
+        elif mode == "match_and_blend":
+            self.synthesized.requires_grad_(False)
+            assert 'criteria' in ext.keys(), 'Please specify a criteria for synthsis.'
+            criteria = make_criteria(ext['criteria']).to(self.device)
+        else:
+            raise ValueError(f'Unsupported mode: {mode}')
+
+        # perform synthsis
+        self.pbar = logger(ext['num_itrs'], len(self.target_pyramid))
+        ext['pbar'] = self.pbar
+        for lvl, lvl_target in enumerate(self.target_pyramid):
+            self.pbar.new_lvl()
+            if lvl > 0:
+                with torch.no_grad():
+                    self.synthesized = F.interpolate(self.synthesized.detach(), size=self.synthesized_lengths[lvl], mode='linear')
+                if mode == "backpropagate":
+                    self.synthesized.requires_grad_(True)
+
+            if mode == "backpropagate": # direct optimize the synthesized motion
+                self.synthesized, losses = GPS.backpropagate(self.synthesized, lvl_target, criteria, ext=ext)
+            elif mode == "match_and_blend":
+                self.synthesized, losses = GPS.match_and_blend(self.synthesized, lvl_target, criteria, ext=ext)
+
+            criteria.clean_cache()
+            if debug_dir:
+                for itr in range(len(losses)):
+                    writer.add_scalar(f'optimize/losses_lvl{lvl}', losses[itr], itr)
+        self.pbar.pbar.close()
+
+
+        return self.synthesized.detach()
+
+    @staticmethod
+    def backpropagate(synthesized, targets, criteria=None, ext=None):
+        """
+        Minimizes criteria(synthesized, target) for num_steps SGD steps
+        Args:
+            targets (torch.Tensor): Target data.
+            ext (dict): extra configurations.
+        """
+        if criteria is None:
+            assert 'criteria' in ext.keys(), 'Criteria is not set'
+            criteria = make_criteria(ext['criteria']).to(synthesized.device)
+
+        optim = None
+        if 'optimizer' in ext.keys():
+            if ext['optimizer'] == 'Adam':
+                optim = torch.optim.Adam([synthesized], lr=ext['lr'])
+            elif ext['optimizer'] == 'SGD':
+                optim = torch.optim.SGD([synthesized], lr=ext['lr'])
+            elif ext['optimizer'] == 'RMSprop':
+                optim = torch.optim.RMSprop([synthesized], lr=ext['lr'])
+            else:
+                print(f'use default RMSprop optimizer')
+        optim = torch.optim.RMSprop([synthesized], lr=ext['lr']) if optim is None else optim
+        # optim = torch.optim.Adam([synthesized], lr=ext['lr']) if optim is None else optim
+        lr_decay = np.exp(np.log(0.333) / ext['num_itrs'])
+
+        # other constraints
+        trajectory = ext['trajectory'] if 'trajectory' in ext.keys() else None
+
+        losses = []
+        for _i in range(ext['num_itrs']):
+            optim.zero_grad()
+            
+            loss = criteria(synthesized, targets)
+
+            if trajectory is not None: ## velo constrain
+                target_traj = F.interpolate(trajectory, size=synthesized.shape[-1], mode='linear')
+                # target_traj = F.interpolate(trajectory, size=synthesized.shape[-1], mode='linear', align_corners=False)
+                target_velo = ext['pos2velo'](target_traj)
+                
+                velo_mask = [-3, -1]
+                loss += 1 * F.l1_loss(synthesized[:, velo_mask, :], target_velo[:, velo_mask, :])
+
+            loss.backward()
+            optim.step()
+
+            # Update staus
+            losses.append(loss.item())
+            if 'pbar' in ext.keys():
+                ext['pbar'].step()
+                ext['pbar'].print()
+
+        return synthesized, losses
+
+    @staticmethod
+    @torch.no_grad()
+    def match_and_blend(synthesized, targets, criteria, ext):
+        """
+        Minimizes criteria(synthesized, target)
+        Args:
+            targets (torch.Tensor): Target data.
+            ext (dict): extra configurations.
+        """
+        losses = []
+        for _i in range(ext['num_itrs']):
+            if 'parts_list' in ext.keys():
+                def extract_part_motions(motion, parts_list):
+                    part_motions = []
+                    n_frames = motion.shape[-1]
+                    rot, pos = motion[:, :-3, :].reshape(-1, 6, n_frames), motion[:, -3:, :]
+
+                    for part in parts_list:
+                        # part -= 1
+                        part = [i -1 for i in part]
+
+                        # print(part)
+                        if 0 in part:
+                            part_motions += [torch.cat([rot[part].view(1, -1, n_frames), pos.view(1, -1, n_frames)], dim=1)]
+                        else:
+                            part_motions += [rot[part].view(1, -1, n_frames)]
+
+                    return part_motions
+                def combine_part_motions(part_motions, parts_list):
+                    assert len(part_motions) == len(parts_list)
+                    n_frames = part_motions[0].shape[-1]
+                    l = max(list(itertools.chain(*parts_list)))
+                    # print(l, n_frames)
+                    # motion = torch.zeros((1, (l+1)*6 + 3, n_frames), device=part_motions[0].device)
+                    rot = torch.zeros(((l+1), 6, n_frames), device=part_motions[0].device)
+                    pos = torch.zeros((1, 3, n_frames), device=part_motions[0].device)
+                    div_rot = torch.zeros((l+1), device=part_motions[0].device)
+                    div_pos = torch.zeros(1, device=part_motions[0].device)
+
+                    for part_motion, part in zip(part_motions, parts_list):
+                        part = [i -1 for i in part]
+
+                        if 0 in part:
+                            # print(part_motion.shape)
+                            pos += part_motion[:, -3:, :]
+                            div_pos += 1
+                            rot[part] += part_motion[:, :-3, :].view(-1, 6, n_frames)
+                            div_rot[part] += 1
+                        else:
+                            rot[part] += part_motion.view(-1, 6, n_frames)
+                            div_rot[part] += 1
+                            
+                    # print(div_rot, div_pos)
+                    # print(rot.shape)
+                    rot = (rot.permute(1, 2, 0) / div_rot).permute(2, 0, 1)
+                    pos = pos / div_pos
+
+                    return torch.cat([rot.view(1, -1, n_frames), pos.view(1, 3, n_frames)], dim=1)
+
+                # raw_synthesized = synthesized
+                # print(synthesized, synthesized.shape)
+                synthesized_part_motions = extract_part_motions(synthesized, ext['parts_list'])
+                targets_part_motions = [extract_part_motions(target, ext['parts_list']) for target in targets]
+
+                synthesized = []
+                for _j in range(len(synthesized_part_motions)):
+                    synthesized_part_motion = synthesized_part_motions[_j]
+                    # synthesized += [synthesized_part_motion]
+                    targets_part_motion = [target[_j] for target in targets_part_motions]
+                    # # print(synthesized_part_motion.shape, targets_part_motion[0].shape)
+                    synthesized += [criteria(synthesized_part_motion, targets_part_motion, ext=ext, return_blended_results=True)[0]]
+
+                # print(len(synthesized))
+                
+                synthesized = combine_part_motions(synthesized, ext['parts_list'])
+                # print(synthesized, synthesized.shape)
+                # print((raw_synthesized-synthesized > 0.00001).sum())
+                # exit()
+                # print(synthesized.shape)
+                losses = 0
+
+                # exit()
+       
+            else:
+                synthesized, loss = criteria(synthesized, targets, ext=ext, return_blended_results=True)
+
+                # Update staus
+                losses.append(loss.item())
+                if 'pbar' in ext.keys():
+                    ext['pbar'].step()
+                    ext['pbar'].print()
+
+        return synthesized, losses
+

NN/losses.py

@@ -0,0 +1,51 @@
+import torch    
+import torch.nn as nn
+
+from .utils import extract_patches, combine_patches, efficient_cdist, get_NNs_Dists
+
+def make_criteria(conf):
+    if conf['type'] == 'PatchCoherentLoss':
+        return PatchCoherentLoss(conf['patch_size'], stride=conf['stride'], loop=conf['loop'], coherent_alpha=conf['coherent_alpha'])
+    elif conf['type'] == 'SWDLoss':
+        raise NotImplementedError('SWDLoss is not implemented')
+    else:
+        raise ValueError('Invalid criteria: {}'.format(conf['criteria']))
+
+class PatchCoherentLoss(torch.nn.Module):
+    def __init__(self, patch_size=7, stride=1, loop=False, coherent_alpha=None, cache=False):
+        super(PatchCoherentLoss, self).__init__()
+        self.patch_size = patch_size
+        self.stride = stride
+        self.loop = loop
+        self.coherent_alpha = coherent_alpha
+        assert self.stride == 1, "Only support stride of 1"
+        # assert self.patch_size % 2 == 1, "Only support odd patch size"
+        self.cache = cache
+        if cache:
+            self.cached_data = None
+
+    def forward(self, X, Ys, dist_wrapper=None, ext=None, return_blended_results=False):
+        """For each patch in input X find its NN in target Y and sum the their distances"""
+        assert X.shape[0] == 1, "Only support batch size of 1"
+        dist_fn = lambda X, Y: dist_wrapper(efficient_cdist, X, Y) if dist_wrapper is not None else efficient_cdist(X, Y)
+
+        x_patches = extract_patches(X, self.patch_size, self.stride, loop=self.loop)
+
+        if not self.cache or self.cached_data is None:
+            y_patches = []
+            for y in Ys:
+                y_patches += [extract_patches(y, self.patch_size, self.stride, loop=False)]
+            y_patches = torch.cat(y_patches, dim=1)
+            self.cached_data = y_patches
+        else:
+            y_patches = self.cached_data
+        
+        nnf, dist = get_NNs_Dists(dist_fn, x_patches.squeeze(0), y_patches.squeeze(0), self.coherent_alpha)
+
+        if return_blended_results:
+            return combine_patches(X.shape, y_patches[:, nnf, :], self.patch_size, self.stride, loop=self.loop), dist.mean()
+        else:
+            return dist.mean()
+    
+    def clean_cache(self):
+        self.cached_data = None

NN/utils.py

@@ -0,0 +1,103 @@
+import torch
+import torch.nn.functional as F
+import unfoldNd
+
+def extract_patches(x, patch_size, stride, loop=False):
+    """Extract patches from a motion sequence"""
+    b, c, _t = x.shape
+
+    # manually padding to loop
+    if loop:
+        half = patch_size // 2
+        front, tail = x[:,:,:half], x[:,:,-half:]
+        x = torch.concat([tail, x, front], dim=-1)
+
+    x_patches = unfoldNd.unfoldNd(x, kernel_size=patch_size, stride=stride).transpose(1, 2).reshape(b, -1, c, patch_size)
+    
+    return x_patches.view(b, -1, c * patch_size)
+
+def combine_patches(x_shape, ys, patch_size, stride, loop=False):
+    """Combine motion patches"""
+    # manually handle to loop
+    out_shape = [*x_shape]
+    if loop:
+        padding = patch_size // 2
+        out_shape[-1] = out_shape[-1] + padding * 2
+
+    combined = unfoldNd.foldNd(ys.permute(0, 2, 1), output_size=tuple(out_shape[-1:]), kernel_size=patch_size, stride=stride)
+
+    # normal fold matrix
+    input_ones = torch.ones(tuple(out_shape), dtype=ys.dtype, device=ys.device)
+    divisor = unfoldNd.unfoldNd(input_ones, kernel_size=patch_size, stride=stride)
+    divisor = unfoldNd.foldNd(divisor, output_size=tuple(out_shape[-1:]), kernel_size=patch_size, stride=stride)
+    combined = (combined / divisor).squeeze(dim=0).unsqueeze(0)
+    
+    if loop:
+        half = patch_size // 2
+        front, tail = combined[:,:,:half], combined[:,:,-half:]
+        combined[:, :, half:2 * half] = (combined[:, :, half:2 * half] + tail) / 2
+        combined[:, :, - 2 * half:-half] = (front + combined[:, :, - 2 * half:-half]) / 2
+        combined = combined[:, :, half:-half]
+
+    return combined
+
+
+def efficient_cdist(X, Y):
+    """
+    Pytorch efficient way of computing distances between all vectors in X and Y, i.e (X[:, None] - Y[None, :])**2
+    Get the nearest neighbor index from Y for each X
+    :param X:  (n1, d) tensor
+    :param Y:  (n2, d) tensor
+    Returns a n2 n1 of indices
+    """
+    dist = (X * X).sum(1)[:, None] + (Y * Y).sum(1)[None, :] - 2.0 * torch.mm(X, torch.transpose(Y, 0, 1))
+    d = X.shape[1]
+    dist /= d # normalize by size of vector to make dists independent of the size of d ( use same alpha for all patche-sizes)
+    return dist # DO NOT use torch.sqrt
+
+
+def get_col_mins_efficient(dist_fn, X, Y, b=1024):
+    """
+    Computes the l2 distance to the closest x or each y.
+    :param X:  (n1, d) tensor
+    :param Y:  (n2, d) tensor
+    Returns n1 long array of L2 distances
+    """
+    n_batches = len(Y) // b
+    mins = torch.zeros(Y.shape[0], dtype=X.dtype, device=X.device)
+    for i in range(n_batches):
+        mins[i * b:(i + 1) * b] = dist_fn(X, Y[i * b:(i + 1) * b]).min(0)[0]
+    if len(Y) % b != 0:
+        mins[n_batches * b:] = dist_fn(X, Y[n_batches * b:]).min(0)[0]
+
+    return mins
+
+
+def get_NNs_Dists(dist_fn, X, Y, alpha=None, b=1024):
+    """
+    Get the nearest neighbor index from Y for each X.
+    Avoids holding a (n1 * n2) amtrix in order to reducing memory footprint to (b * max(n1,n2)).
+    :param X:  (n1, d) tensor
+    :param Y:  (n2, d) tensor
+    Returns a n2 n1 of indices amd distances
+    """
+    if alpha is not None:
+        normalizing_row = get_col_mins_efficient(dist_fn, X, Y, b=b)
+        normalizing_row = alpha + normalizing_row[None, :]
+    else:
+        normalizing_row = 1
+
+    NNs = torch.zeros(X.shape[0], dtype=torch.long, device=X.device)
+    Dists = torch.zeros(X.shape[0], dtype=torch.float, device=X.device)
+
+    n_batches = len(X) // b
+    for i in range(n_batches):
+        dists = dist_fn(X[i * b:(i + 1) * b], Y) / normalizing_row
+        NNs[i * b:(i + 1) * b] = dists.min(1)[1]
+        Dists[i * b:(i + 1) * b] = dists.min(1)[0]
+    if len(X) % b != 0:
+        dists = dist_fn(X[n_batches * b:], Y) / normalizing_row
+        NNs[n_batches * b:] = dists.min(1)[1]
+        Dists[n_batches * b: ] = dists.min(1)[0]
+
+    return NNs, Dists

README.md

@@ -1,8 +1,8 @@
 ---
-title: GenMM Test
-emoji: 😻
-colorFrom: indigo
-colorTo: yellow
+title: GenMM
+emoji: 🌍
+colorFrom: purple
+colorTo: red
 sdk: gradio
 sdk_version: 3.33.1
 app_file: app.py

app.py

@@ -0,0 +1,54 @@
+import json
+import time
+
+from dataset.tracks_motion import TracksMotion
+from GPS import GPS
+import gradio as gr
+
+def _synthesis(synthesis_setting, motion_data):
+    model = GPS(
+        init_mode = f"random_synthesis/{synthesis_setting['frames']}",
+        noise_sigma = synthesis_setting['noise_sigma'],
+        coarse_ratio = 0.2,
+        pyr_factor = synthesis_setting['pyr_factor'],
+        num_stages_limit = -1,
+        silent=True,
+        device='cpu'
+    )
+
+    synthesized_motion = model.run(
+        motion_data, 
+        mode="match_and_blend",
+        ext={
+            'criteria': {
+                'type': 'PatchCoherentLoss',
+                'patch_size': synthesis_setting['patch_size'],
+                'stride': synthesis_setting['stride'] if 'stride' in synthesis_setting.keys() else 1,
+                'loop': synthesis_setting['loop'],
+                'coherent_alpha': synthesis_setting['alpha'] if synthesis_setting['completeness'] else None,
+                },
+            'optimizer': "match_and_blend",
+            'num_itrs': synthesis_setting['num_steps'],
+        }
+    )
+
+    return synthesized_motion
+
+def synthesis(data):
+    data = json.loads(data)
+    # create track object
+    data['setting']['coarse_ratio'] = -1
+    motion_data = TracksMotion(data['tracks'], scale=data['scale'])
+    start = time.time()
+    synthesized_motion = _synthesis(
+        data['setting'], 
+        [motion_data]
+        )
+    end = time.time()
+    data['time'] = end - start
+    data['tracks'] = motion_data.parse(synthesized_motion)
+
+    return data
+
+demo = gr.Interface(fn=synthesis, inputs="json", outputs="json")
+demo.launch()

configs/random_synthesis.yaml

@@ -0,0 +1,36 @@
+outout_dir: './output/random_synthesis'
+
+# for GANimator BVH data
+skeleton_aware: true
+use_velo: true
+repr: 'repr6d'
+contact: true
+keep_y_pos: true
+joint_reduction: true
+
+
+# for synthesis
+coarse_ratio: -1
+coarse_ratio_factor: 10
+pyr_factor: 0.75
+num_stages_limit: -1
+noise_sigma: 10.0
+patch_size: 11
+loop: false
+loss_type: 'PatchCoherent'
+coherent_alpha: 0.01
+optimizer: 'RMSprop'
+lr: 0.01
+num_steps: 3
+decay_rate: 0.9
+decay_steps: 0.9
+
+# for visualization (only for blender render)
+visualization: true
+fbx_path: null
+reso: '[1920, 1080]'
+samples: 64
+fps: 30
+frame_end: -1
+camera_pos: '[0, -8, 2.5]'
+target_pos: '[0, 2, 0.5]'

dataset/tracks_motion.py

@@ -0,0 +1,183 @@
+import os
+from os.path import join as pjoin
+import numpy as np
+import copy
+import torch
+import torch.nn.functional as F
+from utils.transforms import quat2repr6d, quat2euler, repr6d2quat
+
+class TracksParser():
+    def __init__(self, tracks_json, scale=1.0, requires_contact=False, joint_reduction=False):
+        assert requires_contact==False, 'contact is not implemented for tracks data yet!!!'
+
+        self.tracks_json = tracks_json
+        self.scale = scale
+        self.requires_contact = requires_contact
+        self.joint_reduction = joint_reduction
+        
+        self.skeleton_names = []
+        self.rotations = []
+        for i, track in enumerate(self.tracks_json):
+            # print(i, track['name'])
+            self.skeleton_names.append(track['name'])
+            if i == 0:
+                assert track['type'] == 'vector'
+                self.position = np.array(track['values']).reshape(-1, 3) * self.scale
+                self.num_frames = self.position.shape[0]
+            else:
+                assert track['type'] == 'quaternion' # DEAFULT: quaternion
+                rotation = np.array(track['values']).reshape(-1, 4)
+                if rotation.shape[0] == 0:
+                    rotation = np.zeros((self.num_frames, 4))
+                elif rotation.shape[0] < self.num_frames:
+                    rotation = np.repeat(rotation, self.num_frames // rotation.shape[0], axis=0)
+                elif rotation.shape[0] > self.num_frames:
+                    rotation = rotation[:self.num_frames]
+                self.rotations += [rotation]
+        self.rotations = np.array(self.rotations, dtype=np.float32)
+
+    def to_tensor(self, repr='euler', rot_only=False):
+        if repr not in ['euler', 'quat', 'quaternion', 'repr6d']:
+            raise Exception('Unknown rotation representation')
+        rotations = self.get_rotation(repr=repr)
+        positions = self.get_position()
+
+        if rot_only:
+            return rotations.reshape(rotations.shape[0], -1)
+
+        if self.requires_contact:
+            virtual_contact = torch.zeros_like(rotations[:, :len(self.skeleton.contact_id)])
+            virtual_contact[..., 0] = self.contact_label
+            rotations = torch.cat([rotations, virtual_contact], dim=1)
+
+        rotations = rotations.reshape(rotations.shape[0], -1)
+        return torch.cat((rotations, positions), dim=-1)
+
+    def get_rotation(self, repr='quat'):
+        if repr == 'quaternion' or repr == 'quat' or repr == 'repr6d':
+            rotations = torch.tensor(self.rotations, dtype=torch.float).transpose(0, 1)
+        if repr == 'repr6d':
+            rotations = quat2repr6d(rotations)
+        if repr == 'euler':
+            rotations = quat2euler(rotations)
+        return rotations
+
+    def get_position(self):
+        return torch.tensor(self.position, dtype=torch.float32)
+
+class TracksMotion:
+    def __init__(self, tracks_json, scale=1.0, repr='repr6d', padding=False,
+                 use_velo=True, contact=False, keep_y_pos=True, joint_reduction=False):
+        self.scale = scale
+        self.tracks = TracksParser(tracks_json, scale, requires_contact=contact, joint_reduction=joint_reduction)
+        self.raw_motion = self.tracks.to_tensor(repr=repr)
+        self.extra = {
+
+        }
+
+        self.repr = repr
+        if repr == 'quat':
+            self.n_rot = 4
+        elif repr == 'repr6d':
+            self.n_rot = 6
+        elif repr == 'euler':
+            self.n_rot = 3
+        self.padding = padding
+        self.use_velo = use_velo
+        self.contact = contact
+        self.keep_y_pos = keep_y_pos
+        self.joint_reduction = joint_reduction
+
+        self.raw_motion = self.raw_motion.permute(1, 0).unsqueeze_(0) # Shape = (1, n_channel, n_frames)
+        self.extra['global_pos'] = self.raw_motion[:, -3:, :]
+
+        if padding:
+            self.n_pad = self.n_rot - 3 # pad position channels
+            paddings = torch.zeros_like(self.raw_motion[:, :self.n_pad])
+            self.raw_motion = torch.cat((self.raw_motion, paddings), dim=1)
+        else:
+            self.n_pad = 0
+        self.raw_motion = torch.cat((self.raw_motion[:, :-3-self.n_pad], self.raw_motion[:, -3-self.n_pad:]), dim=1)
+
+        if self.use_velo:
+            self.msk = [-3, -2, -1] if not keep_y_pos else [-3, -1]
+            self.raw_motion = self.pos2velo(self.raw_motion)
+
+        self.n_contact = len(self.tracks.skeleton.contact_id) if contact else 0
+
+    @property
+    def n_channels(self):
+        return self.raw_motion.shape[1]
+
+    def __len__(self):
+        return self.raw_motion.shape[-1]
+
+    def pos2velo(self, pos):
+        msk = [i - self.n_pad for i in self.msk]
+        velo = pos.detach().clone().to(pos.device)
+        velo[:, msk, 1:] = pos[:, msk, 1:] - pos[:, msk, :-1]
+        self.begin_pos = pos[:, msk, 0].clone()
+        velo[:, msk, 0] = pos[:, msk, 1]
+        return velo
+
+    def velo2pos(self, velo):
+        msk = [i - self.n_pad for i in self.msk]
+        pos = velo.detach().clone().to(velo.device)
+        pos[:, msk, 0] = self.begin_pos.to(velo.device)
+        pos[:, msk] = torch.cumsum(velo[:, msk], dim=-1)
+        return pos
+
+    def motion2pos(self, motion):
+        if not self.use_velo:
+            return motion
+        else:
+            self.velo2pos(motion.clone())
+
+    def sample(self, size=None, slerp=False, align_corners=False):
+        if size is None:
+            return {'motion': self.raw_motion, 'extra': self.extra}
+        else:
+            if slerp:
+                raise NotImplementedError('slerp is not not implemented yet!!!')
+            else:
+                motion = F.interpolate(self.raw_motion, size=size, mode='linear', align_corners=align_corners)
+                extra = {}
+                if 'global_pos' in self.extra.keys():
+                    extra['global_pos'] = F.interpolate(self.extra['global_pos'], size=size, mode='linear', align_corners=align_corners)
+
+            return motion
+            # return {'motion': motion, 'extra': extra}
+
+    def parse(self, motion, keep_velo=False,):
+        """
+        No batch support here!!!
+        :returns tracks_json
+        """
+        motion = motion.clone()
+
+        if self.use_velo and not keep_velo:
+            motion = self.velo2pos(motion)
+        if self.n_pad:
+            motion = motion[:, :-self.n_pad]
+        if self.contact:
+            raise NotImplementedError('contact is not implemented yet!!!')
+
+        motion = motion.squeeze().permute(1, 0)
+        pos = motion[..., -3:] / self.scale
+        rot = motion[..., :-3].reshape(motion.shape[0], -1, self.n_rot)
+        if self.repr == 'repr6d':
+            rot = repr6d2quat(rot)
+        elif self.repr == 'euler':
+            raise NotImplementedError('parse "euler is not implemented yet!!!')
+
+        times = []
+        out_tracks_json = copy.deepcopy(self.tracks.tracks_json)
+        for i, _track in enumerate(out_tracks_json):
+            if i == 0:
+                times = [ j * out_tracks_json[i]['times'][1] for j in range(motion.shape[0])]
+                out_tracks_json[i]['values'] = pos.flatten().detach().cpu().numpy().tolist() 
+            else:
+                out_tracks_json[i]['values'] = rot[:, i-1, :].flatten().detach().cpu().numpy().tolist()
+            out_tracks_json[i]['times'] = times
+
+        return out_tracks_json

requirements.txt

@@ -0,0 +1,15 @@
+filterpy==1.4.5
+torchvision==0.12.0
+tensorboardX==2.5
+protobuf==3.20.1
+scipy==1.7.3
+tqdm==4.62.3
+unfoldNd
+flask==2.1.3
+flask-cors==3.0.10
+pyyaml>=5.3.1
+requests
+tensorboard
+transforms3d
+imageio
+imageio-ffmpeg

utils/base.py

@@ -0,0 +1,148 @@
+import os
+import os.path as osp
+import sys
+import time
+import yaml
+import imageio
+import random
+import shutil
+import random
+import numpy as np
+import torch
+from tqdm import tqdm
+import matplotlib.pyplot as plt
+
+class ConfigParser():
+    def __init__(self, args):
+        """
+        class to parse configuration.
+        """
+        args = args.parse_args()
+        self.cfg = self.merge_config_file(args)
+
+        # set random seed
+        self.set_seed()
+
+    def __str__(self):
+        return str(self.cfg.__dict__)
+
+    def __getattr__(self, name):
+        """
+        Access items use dot.notation.
+        """
+        return self.cfg.__dict__[name]
+
+    def __getitem__(self, name):
+        """
+        Access items like ordinary dict.
+        """
+        return self.cfg.__dict__[name]
+
+    def merge_config_file(self, args, allow_invalid=True):
+        """
+        Load json config file and merge the arguments
+        """
+        assert args.config is not None
+        with open(args.config, 'r') as f:
+            cfg = yaml.safe_load(f)
+            if 'config' in cfg.keys():
+                del cfg['config']
+        f.close()
+        invalid_args = list(set(cfg.keys()) - set(dir(args)))
+        if invalid_args and not allow_invalid:
+            raise ValueError(f"Invalid args {invalid_args} in {args.config}.")
+        
+        for k in list(cfg.keys()):
+            if k in args.__dict__.keys() and args.__dict__[k] is not None:
+                print('=========>  overwrite config: {} = {}'.format(k, args.__dict__[k]))
+                del cfg[k]
+
+        args.__dict__.update(cfg)
+
+        return args
+
+    def set_seed(self):
+        ''' set random seed for random, numpy and torch. '''
+        if 'seed' not in self.cfg.__dict__.keys():
+            return
+        if self.cfg.seed is None:
+            self.cfg.seed = int(time.time()) % 1000000
+        print('=========>  set random seed: {}'.format(self.cfg.seed))
+        # fix random seeds for reproducibility
+        random.seed(self.cfg.seed)
+        np.random.seed(self.cfg.seed)
+        torch.manual_seed(self.cfg.seed)
+        torch.cuda.manual_seed(self.cfg.seed)
+
+    def save_codes_and_config(self, save_path):
+        """
+        save codes and config to $save_path.
+        """
+        cur_codes_path = osp.dirname(osp.dirname(os.path.abspath(__file__)))
+        if os.path.exists(save_path):
+            shutil.rmtree(save_path)
+        shutil.copytree(cur_codes_path, osp.join(save_path, 'codes'), \
+            ignore=shutil.ignore_patterns('*debug*', '*data*', '*output*', '*exps*', '*.txt', '*.json', '*.mp4', '*.png', '*.jpg', '*.bvh', '*.csv', '*.pth', '*.tar', '*.npz'))
+
+        with open(osp.join(save_path, 'config.yaml'), 'w') as f:
+            f.write(yaml.dump(self.cfg.__dict__))
+        f.close()
+
+
+# other utils
+class logger:
+    """Keeps track of the levels and steps of optimization. Logs it via TQDM"""
+    def __init__(self, n_steps, n_lvls):
+        self.n_steps = n_steps
+        self.n_lvls = n_lvls
+        self.lvl = -1
+        self.lvl_step = 0
+        self.steps = 0
+        self.pbar = tqdm(total=self.n_lvls * self.n_steps, desc='Starting')
+
+    def step(self):
+        self.pbar.update(1)
+        self.steps += 1
+        self.lvl_step += 1
+
+    def new_lvl(self):
+        self.lvl += 1
+        self.lvl_step = 0
+
+    def print(self):
+        self.pbar.set_description(f'Lvl {self.lvl}/{self.n_lvls-1}, step {self.lvl_step}/{self.n_steps}')
+
+
+def set_seed(seed):
+    if seed is not None:
+        random.seed(seed)
+        np.random.seed(seed)
+        torch.manual_seed(seed)
+        torch.cuda.manual_seed(seed)
+
+
+# debug utils
+def draw_trajectory(trajectory, save_path=None, anim=True):
+    r = max(abs(trajectory.min()), trajectory.max())
+    if anim:
+        imgs = []
+        for i in tqdm(range(1, trajectory.shape[0])):
+            plt.plot(trajectory[:i, 0], trajectory[:i, 2], color='red')
+            plt.xlim(-r-1, r+1)
+            plt.ylim(-r-1, r+1)
+            plt.savefig(save_path + '.png')
+            imgs += [imageio.imread(save_path + '.png')]
+        imageio.mimwrite(save_path + '.mp4', imgs)
+        plt.close()
+    else:
+        # plt.scatter(trajectory[:, 0], trajectory[:, 1], trajectory[:, 2])
+        plt.plot(trajectory[:, 0], trajectory[:, 2], color='red')
+        plt.xlim(-r*1.5, r*1.5)
+        plt.ylim(-r*1.5, r*1.5)
+        if save_path is not None:
+            plt.savefig(save_path + '.png')
+            plt.close()
+
+    # velo = self.raw_motion[0, self.mask, :].numpy()
+    # print(velo.shape)
+    # imgs = []

utils/contact.py

@@ -0,0 +1,103 @@
+import torch
+
+
+def foot_contact_by_height(pos):
+    eps = 0.25
+    return (-eps < pos[..., 1]) * (pos[..., 1] < eps)
+
+
+def velocity(pos, padding=False):
+    velo = pos[1:, ...] - pos[:-1, ...]
+    velo_norm = torch.norm(velo, dim=-1)
+    if padding:
+        pad = torch.zeros_like(velo_norm[:1, :])
+        velo_norm = torch.cat([pad, velo_norm], dim=0)
+    return velo_norm
+
+
+def foot_contact(pos, ref_height=1., threshold=0.018):
+    velo_norm = velocity(pos)
+    contact = velo_norm < threshold
+    contact = contact.int()
+    padding = torch.zeros_like(contact)
+    contact = torch.cat([padding[:1, :], contact], dim=0)
+    return contact
+
+
+def alpha(t):
+    return 2.0 * t * t * t - 3.0 * t * t + 1
+
+
+def lerp(a, l, r):
+    return (1 - a) * l + a * r
+
+
+def constrain_from_contact(contact, glb, fid='TBD', L=5):
+    """
+    :param contact: contact label
+    :param glb: original global position
+    :param fid: joint id to fix, corresponding to the order in contact
+    :param L: frame to look forward/backward
+    :return:
+    """
+    T = glb.shape[0]
+
+    for i, fidx in enumerate(fid):  # fidx: index of the foot joint
+        fixed = contact[:, i]  # [T]
+        s = 0
+        while s < T:
+            while s < T and fixed[s] == 0:
+                s += 1
+            if s >= T:
+                break
+            t = s
+            avg = glb[t, fidx].clone()
+            while t + 1 < T and fixed[t + 1] == 1:
+                t += 1
+                avg += glb[t, fidx].clone()
+            avg /= (t - s + 1)
+
+            for j in range(s, t + 1):
+                glb[j, fidx] = avg.clone()
+            s = t + 1
+
+        for s in range(T):
+            if fixed[s] == 1:
+                continue
+            l, r = None, None
+            consl, consr = False, False
+            for k in range(L):
+                if s - k - 1 < 0:
+                    break
+                if fixed[s - k - 1]:
+                    l = s - k - 1
+                    consl = True
+                    break
+            for k in range(L):
+                if s + k + 1 >= T:
+                    break
+                if fixed[s + k + 1]:
+                    r = s + k + 1
+                    consr = True
+                    break
+            if not consl and not consr:
+                continue
+            if consl and consr:
+                litp = lerp(alpha(1.0 * (s - l + 1) / (L + 1)),
+                            glb[s, fidx], glb[l, fidx])
+                ritp = lerp(alpha(1.0 * (r - s + 1) / (L + 1)),
+                            glb[s, fidx], glb[r, fidx])
+                itp = lerp(alpha(1.0 * (s - l + 1) / (r - l + 1)),
+                           ritp, litp)
+                glb[s, fidx] = itp.clone()
+                continue
+            if consl:
+                litp = lerp(alpha(1.0 * (s - l + 1) / (L + 1)),
+                            glb[s, fidx], glb[l, fidx])
+                glb[s, fidx] = litp.clone()
+                continue
+            if consr:
+                ritp = lerp(alpha(1.0 * (r - s + 1) / (L + 1)),
+                            glb[s, fidx], glb[r, fidx])
+                glb[s, fidx] = ritp.clone()
+    return glb

utils/kinematics.py

@@ -0,0 +1,203 @@
+import torch
+from utils.transforms import quat2mat, repr6d2mat, euler2mat
+
+
+class ForwardKinematics:
+    def __init__(self, parents, offsets=None):
+        self.parents = parents
+        if offsets is not None and len(offsets.shape) == 2:
+            offsets = offsets.unsqueeze(0)
+        self.offsets = offsets
+
+    def forward(self, rots, offsets=None, global_pos=None):
+        """
+        Forward Kinematics: returns a per-bone transformation
+        @param rots: local joint rotations (batch_size, bone_num, 3, 3)
+        @param offsets: (batch_size, bone_num, 3) or None
+        @param global_pos: global_position: (batch_size, 3) or keep it as in offsets (default)
+        @return: (batch_szie, bone_num, 3, 4)
+        """
+        rots = rots.clone()
+        if offsets is None:
+            offsets = self.offsets.to(rots.device)
+        if global_pos is None:
+            global_pos = offsets[:, 0]
+
+        pos = torch.zeros((rots.shape[0], rots.shape[1], 3), device=rots.device)
+        rest_pos = torch.zeros_like(pos)
+        res = torch.zeros((rots.shape[0], rots.shape[1], 3, 4), device=rots.device)
+
+        pos[:, 0] = global_pos
+        rest_pos[:, 0] = offsets[:, 0]
+
+        for i, p in enumerate(self.parents):
+            if i != 0:
+                rots[:, i] = torch.matmul(rots[:, p], rots[:, i])
+                pos[:, i] = torch.matmul(rots[:, p], offsets[:, i].unsqueeze(-1)).squeeze(-1) + pos[:, p]
+                rest_pos[:, i] = rest_pos[:, p] + offsets[:, i]
+
+            res[:, i, :3, :3] = rots[:, i]
+            res[:, i, :, 3] = torch.matmul(rots[:, i], -rest_pos[:, i].unsqueeze(-1)).squeeze(-1) + pos[:, i]
+
+        return res
+
+    def accumulate(self, local_rots):
+        """
+        Get global joint rotation from local rotations
+        @param local_rots: (batch_size, n_bone, 3, 3)
+        @return: global_rotations
+        """
+        res = torch.empty_like(local_rots)
+        for i, p in enumerate(self.parents):
+            if i == 0:
+                res[:, i] = local_rots[:, i]
+            else:
+                res[:, i] = torch.matmul(res[:, p], local_rots[:, i])
+        return res
+
+    def unaccumulate(self, global_rots):
+        """
+        Get local joint rotation from global rotations
+        @param global_rots: (batch_size, n_bone, 3, 3)
+        @return: local_rotations
+        """
+        res = torch.empty_like(global_rots)
+        inv = torch.empty_like(global_rots)
+
+        for i, p in enumerate(self.parents):
+            if i == 0:
+                inv[:, i] = global_rots[:, i].transpose(-2, -1)
+                res[:, i] = global_rots[:, i]
+                continue
+            res[:, i] = torch.matmul(inv[:, p], global_rots[:, i])
+            inv[:, i] = torch.matmul(res[:, i].transpose(-2, -1), inv[:, p])
+
+        return res
+
+
+class ForwardKinematicsJoint:
+    def __init__(self, parents, offset):
+        self.parents = parents
+        self.offset = offset
+
+    '''
+        rotation should have shape batch_size * Joint_num * (3/4) * Time
+        position should have shape batch_size * 3 * Time
+        offset should have shape batch_size * Joint_num * 3
+        output have shape batch_size * Time * Joint_num * 3
+    '''
+
+    def forward(self, rotation: torch.Tensor, position: torch.Tensor, offset=None,
+                world=True):
+        '''
+        if not quater and rotation.shape[-2] != 3: raise Exception('Unexpected shape of rotation')
+        if quater and rotation.shape[-2] != 4: raise Exception('Unexpected shape of rotation')
+        rotation = rotation.permute(0, 3, 1, 2)
+        position = position.permute(0, 2, 1)
+        '''
+        if rotation.shape[-1] == 6:
+            transform = repr6d2mat(rotation)
+        elif rotation.shape[-1] == 4:
+            norm = torch.norm(rotation, dim=-1, keepdim=True)
+            rotation = rotation / norm
+            transform = quat2mat(rotation)
+        elif rotation.shape[-1] == 3:
+            transform = euler2mat(rotation)
+        else:
+            raise Exception('Only accept quaternion rotation input')
+        result = torch.empty(transform.shape[:-2] + (3,), device=position.device)
+
+        if offset is None:
+            offset = self.offset
+        offset = offset.reshape((-1, 1, offset.shape[-2], offset.shape[-1], 1))
+
+        result[..., 0, :] = position
+        for i, pi in enumerate(self.parents):
+            if pi == -1:
+                assert i == 0
+                continue
+
+            result[..., i, :] = torch.matmul(transform[..., pi, :, :], offset[..., i, :, :]).squeeze()
+            transform[..., i, :, :] = torch.matmul(transform[..., pi, :, :].clone(), transform[..., i, :, :].clone())
+            if world: result[..., i, :] += result[..., pi, :]
+        return result
+
+
+class InverseKinematicsJoint:
+    def __init__(self, rotations: torch.Tensor, positions: torch.Tensor, offset, parents, constrains):
+        self.rotations = rotations.detach().clone()
+        self.rotations.requires_grad_(True)
+        self.position = positions.detach().clone()
+        self.position.requires_grad_(True)
+
+        self.parents = parents
+        self.offset = offset
+        self.constrains = constrains
+
+        self.optimizer = torch.optim.Adam([self.position, self.rotations], lr=1e-3, betas=(0.9, 0.999))
+        self.criteria = torch.nn.MSELoss()
+
+        self.fk = ForwardKinematicsJoint(parents, offset)
+
+        self.glb = None
+
+    def step(self):
+        self.optimizer.zero_grad()
+        glb = self.fk.forward(self.rotations, self.position)
+        loss = self.criteria(glb, self.constrains)
+        loss.backward()
+        self.optimizer.step()
+        self.glb = glb
+        return loss.item()
+
+
+class InverseKinematicsJoint2:
+    def __init__(self, rotations: torch.Tensor, positions: torch.Tensor, offset, parents, constrains, cid,
+                 lambda_rec_rot=1., lambda_rec_pos=1., use_velo=False):
+        self.use_velo = use_velo
+        self.rotations_ori = rotations.detach().clone()
+        self.rotations = rotations.detach().clone()
+        self.rotations.requires_grad_(True)
+        self.position_ori = positions.detach().clone()
+        self.position = positions.detach().clone()
+        if self.use_velo:
+            self.position[1:] = self.position[1:] - self.position[:-1]
+        self.position.requires_grad_(True)
+
+        self.parents = parents
+        self.offset = offset
+        self.constrains = constrains.detach().clone()
+        self.cid = cid
+
+        self.lambda_rec_rot = lambda_rec_rot
+        self.lambda_rec_pos = lambda_rec_pos
+
+        self.optimizer = torch.optim.Adam([self.position, self.rotations], lr=1e-3, betas=(0.9, 0.999))
+        self.criteria = torch.nn.MSELoss()
+
+        self.fk = ForwardKinematicsJoint(parents, offset)
+
+        self.glb = None
+
+    def step(self):
+        self.optimizer.zero_grad()
+        if self.use_velo:
+            position = torch.cumsum(self.position, dim=0)
+        else:
+            position = self.position
+        glb = self.fk.forward(self.rotations, position)
+        self.constrain_loss = self.criteria(glb[:, self.cid], self.constrains)
+        self.rec_loss_rot = self.criteria(self.rotations, self.rotations_ori)
+        self.rec_loss_pos = self.criteria(self.position, self.position_ori)
+        loss = self.constrain_loss + self.rec_loss_rot * self.lambda_rec_rot + self.rec_loss_pos * self.lambda_rec_pos
+        loss.backward()
+        self.optimizer.step()
+        self.glb = glb
+        return loss.item()
+
+    def get_position(self):
+        if self.use_velo:
+            position = torch.cumsum(self.position.detach(), dim=0)
+        else:
+            position = self.position.detach()
+        return position

utils/skeleton.py

@@ -0,0 +1,347 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+import numpy as np
+
+
+class SkeletonConv(nn.Module):
+    def __init__(self, neighbour_list, in_channels, out_channels, kernel_size, joint_num, stride=1, padding=0,
+                 bias=True, padding_mode='zeros', add_offset=False, in_offset_channel=0):
+        super(SkeletonConv, self).__init__()
+
+        if in_channels % joint_num != 0 or out_channels % joint_num != 0:
+            raise Exception('in/out channels should be divided by joint_num')
+        self.in_channels_per_joint = in_channels // joint_num
+        self.out_channels_per_joint = out_channels // joint_num
+
+        if padding_mode == 'zeros': padding_mode = 'constant'
+
+        self.expanded_neighbour_list = []
+        self.expanded_neighbour_list_offset = []
+        self.neighbour_list = neighbour_list
+        self.add_offset = add_offset
+        self.joint_num = joint_num
+
+        self.stride = stride
+        self.dilation = 1
+        self.groups = 1
+        self.padding = padding
+        self.padding_mode = padding_mode
+        self._padding_repeated_twice = (padding, padding)
+
+        for neighbour in neighbour_list:
+            expanded = []
+            for k in neighbour:
+                for i in range(self.in_channels_per_joint):
+                    expanded.append(k * self.in_channels_per_joint + i)
+            self.expanded_neighbour_list.append(expanded)
+
+        if self.add_offset:
+            self.offset_enc = SkeletonLinear(neighbour_list, in_offset_channel * len(neighbour_list), out_channels)
+
+            for neighbour in neighbour_list:
+                expanded = []
+                for k in neighbour:
+                    for i in range(add_offset):
+                        expanded.append(k * in_offset_channel + i)
+                self.expanded_neighbour_list_offset.append(expanded)
+
+        self.weight = torch.zeros(out_channels, in_channels, kernel_size)
+        if bias:
+            self.bias = torch.zeros(out_channels)
+        else:
+            self.register_parameter('bias', None)
+
+        self.mask = torch.zeros_like(self.weight)
+        for i, neighbour in enumerate(self.expanded_neighbour_list):
+            self.mask[self.out_channels_per_joint * i: self.out_channels_per_joint * (i + 1), neighbour, ...] = 1
+        self.mask = nn.Parameter(self.mask, requires_grad=False)
+
+        self.description = 'SkeletonConv(in_channels_per_armature={}, out_channels_per_armature={}, kernel_size={}, ' \
+                           'joint_num={}, stride={}, padding={}, bias={})'.format(
+            in_channels // joint_num, out_channels // joint_num, kernel_size, joint_num, stride, padding, bias
+        )
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        for i, neighbour in enumerate(self.expanded_neighbour_list):
+            """ Use temporary variable to avoid assign to copy of slice, which might lead to un expected result """
+            tmp = torch.zeros_like(self.weight[self.out_channels_per_joint * i: self.out_channels_per_joint * (i + 1),
+                                   neighbour, ...])
+            nn.init.kaiming_uniform_(tmp, a=math.sqrt(5))
+            self.weight[self.out_channels_per_joint * i: self.out_channels_per_joint * (i + 1),
+                        neighbour, ...] = tmp
+            if self.bias is not None:
+                fan_in, _ = nn.init._calculate_fan_in_and_fan_out(
+                    self.weight[self.out_channels_per_joint * i: self.out_channels_per_joint * (i + 1), neighbour, ...])
+                bound = 1 / math.sqrt(fan_in)
+                tmp = torch.zeros_like(
+                    self.bias[self.out_channels_per_joint * i: self.out_channels_per_joint * (i + 1)])
+                nn.init.uniform_(tmp, -bound, bound)
+                self.bias[self.out_channels_per_joint * i: self.out_channels_per_joint * (i + 1)] = tmp
+
+        self.weight = nn.Parameter(self.weight)
+        if self.bias is not None:
+            self.bias = nn.Parameter(self.bias)
+
+    def set_offset(self, offset):
+        if not self.add_offset: raise Exception('Wrong Combination of Parameters')
+        self.offset = offset.reshape(offset.shape[0], -1)
+
+    def forward(self, input):
+        weight_masked = self.weight * self.mask
+        res = F.conv1d(F.pad(input, self._padding_repeated_twice, mode=self.padding_mode),
+                       weight_masked, self.bias, self.stride,
+                       0, self.dilation, self.groups)
+
+        if self.add_offset:
+            offset_res = self.offset_enc(self.offset)
+            offset_res = offset_res.reshape(offset_res.shape + (1, ))
+            res += offset_res / 100
+        return res
+
+    def __repr__(self):
+        return self.description
+
+
+class SkeletonLinear(nn.Module):
+    def __init__(self, neighbour_list, in_channels, out_channels, extra_dim1=False):
+        super(SkeletonLinear, self).__init__()
+        self.neighbour_list = neighbour_list
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.in_channels_per_joint = in_channels // len(neighbour_list)
+        self.out_channels_per_joint = out_channels // len(neighbour_list)
+        self.extra_dim1 = extra_dim1
+        self.expanded_neighbour_list = []
+
+        for neighbour in neighbour_list:
+            expanded = []
+            for k in neighbour:
+                for i in range(self.in_channels_per_joint):
+                    expanded.append(k * self.in_channels_per_joint + i)
+            self.expanded_neighbour_list.append(expanded)
+
+        self.weight = torch.zeros(out_channels, in_channels)
+        self.mask = torch.zeros(out_channels, in_channels)
+        self.bias = nn.Parameter(torch.Tensor(out_channels))
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        for i, neighbour in enumerate(self.expanded_neighbour_list):
+            tmp = torch.zeros_like(
+                self.weight[i*self.out_channels_per_joint: (i + 1)*self.out_channels_per_joint, neighbour]
+            )
+            self.mask[i*self.out_channels_per_joint: (i + 1)*self.out_channels_per_joint, neighbour] = 1
+            nn.init.kaiming_uniform_(tmp, a=math.sqrt(5))
+            self.weight[i*self.out_channels_per_joint: (i + 1)*self.out_channels_per_joint, neighbour] = tmp
+
+        fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight)
+        bound = 1 / math.sqrt(fan_in)
+        nn.init.uniform_(self.bias, -bound, bound)
+
+        self.weight = nn.Parameter(self.weight)
+        self.mask = nn.Parameter(self.mask, requires_grad=False)
+
+    def forward(self, input):
+        input = input.reshape(input.shape[0], -1)
+        weight_masked = self.weight * self.mask
+        res = F.linear(input, weight_masked, self.bias)
+        if self.extra_dim1: res = res.reshape(res.shape + (1,))
+        return res
+
+
+class SkeletonPoolJoint(nn.Module):
+    def __init__(self, topology, pooling_mode, channels_per_joint, last_pool=False):
+        super(SkeletonPoolJoint, self).__init__()
+
+        if pooling_mode != 'mean':
+            raise Exception('Unimplemented pooling mode in matrix_implementation')
+
+        self.joint_num = len(topology)
+        self.parent = topology
+        self.pooling_list = []
+        self.pooling_mode = pooling_mode
+
+        self.pooling_map = [-1 for _ in range(len(self.parent))]
+        self.child = [-1 for _ in range(len(self.parent))]
+        children_cnt = [0 for _ in range(len(self.parent))]
+        for x, pa in enumerate(self.parent):
+            if pa < 0: continue
+            children_cnt[pa] += 1
+            self.child[pa] = x
+        self.pooling_map[0] = 0
+        for x in range(len(self.parent)):
+            if children_cnt[x] == 0 or (children_cnt[x] == 1 and children_cnt[self.child[x]] > 1):
+                while children_cnt[x] <= 1:
+                    pa = self.parent[x]
+                    if last_pool:
+                        seq = [x]
+                        while pa != -1 and children_cnt[pa] == 1:
+                            seq = [pa] + seq
+                            x = pa
+                            pa = self.parent[x]
+                        self.pooling_list.append(seq)
+                        break
+                    else:
+                        if pa != -1 and children_cnt[pa] == 1:
+                            self.pooling_list.append([pa, x])
+                            x = self.parent[pa]
+                        else:
+                            self.pooling_list.append([x, ])
+                            break
+            elif children_cnt[x] > 1:
+                self.pooling_list.append([x, ])
+
+        self.description = 'SkeletonPool(in_joint_num={}, out_joint_num={})'.format(
+            len(topology), len(self.pooling_list),
+        )
+
+        self.pooling_list.sort(key=lambda x:x[0])
+        for i, a in enumerate(self.pooling_list):
+            for j in a:
+                self.pooling_map[j] = i
+
+        self.output_joint_num = len(self.pooling_list)
+        self.new_topology = [-1 for _ in range(len(self.pooling_list))]
+        for i, x in enumerate(self.pooling_list):
+            if i < 1: continue
+            self.new_topology[i] = self.pooling_map[self.parent[x[0]]]
+
+        self.weight = torch.zeros(len(self.pooling_list) * channels_per_joint, self.joint_num * channels_per_joint)
+
+        for i, pair in enumerate(self.pooling_list):
+            for j in pair:
+                for c in range(channels_per_joint):
+                    self.weight[i * channels_per_joint + c, j * channels_per_joint + c] = 1.0 / len(pair)
+
+        self.weight = nn.Parameter(self.weight, requires_grad=False)
+
+    def forward(self, input: torch.Tensor):
+        return torch.matmul(self.weight, input.unsqueeze(-1)).squeeze(-1)
+
+
+class SkeletonPool(nn.Module):
+    def __init__(self, edges, pooling_mode, channels_per_edge, last_pool=False):
+        super(SkeletonPool, self).__init__()
+
+        if pooling_mode != 'mean':
+            raise Exception('Unimplemented pooling mode in matrix_implementation')
+
+        self.channels_per_edge = channels_per_edge
+        self.pooling_mode = pooling_mode
+        self.edge_num = len(edges) + 1
+        self.seq_list = []
+        self.pooling_list = []
+        self.new_edges = []
+        degree = [0] * 100
+
+        for edge in edges:
+            degree[edge[0]] += 1
+            degree[edge[1]] += 1
+
+        def find_seq(j, seq):
+            nonlocal self, degree, edges
+
+            if degree[j] > 2 and j != 0:
+                self.seq_list.append(seq)
+                seq = []
+
+            if degree[j] == 1:
+                self.seq_list.append(seq)
+                return
+
+            for idx, edge in enumerate(edges):
+                if edge[0] == j:
+                    find_seq(edge[1], seq + [idx])
+
+        find_seq(0, [])
+        for seq in self.seq_list:
+            if last_pool:
+                self.pooling_list.append(seq)
+                continue
+            if len(seq) % 2 == 1:
+                self.pooling_list.append([seq[0]])
+                self.new_edges.append(edges[seq[0]])
+                seq = seq[1:]
+            for i in range(0, len(seq), 2):
+                self.pooling_list.append([seq[i], seq[i + 1]])
+                self.new_edges.append([edges[seq[i]][0], edges[seq[i + 1]][1]])
+
+        # add global position
+        self.pooling_list.append([self.edge_num - 1])
+
+        self.description = 'SkeletonPool(in_edge_num={}, out_edge_num={})'.format(
+            len(edges), len(self.pooling_list)
+        )
+
+        self.weight = torch.zeros(len(self.pooling_list) * channels_per_edge, self.edge_num * channels_per_edge)
+
+        for i, pair in enumerate(self.pooling_list):
+            for j in pair:
+                for c in range(channels_per_edge):
+                    self.weight[i * channels_per_edge + c, j * channels_per_edge + c] = 1.0 / len(pair)
+
+        self.weight = nn.Parameter(self.weight, requires_grad=False)
+
+    def forward(self, input: torch.Tensor):
+        return torch.matmul(self.weight, input)
+
+
+class SkeletonUnpool(nn.Module):
+    def __init__(self, pooling_list, channels_per_edge):
+        super(SkeletonUnpool, self).__init__()
+        self.pooling_list = pooling_list
+        self.input_joint_num = len(pooling_list)
+        self.output_joint_num = 0
+        self.channels_per_edge = channels_per_edge
+        for t in self.pooling_list:
+            self.output_joint_num += len(t)
+
+        self.description = 'SkeletonUnpool(in_joint_num={}, out_joint_num={})'.format(
+            self.input_joint_num, self.output_joint_num,
+        )
+
+        self.weight = torch.zeros(self.output_joint_num * channels_per_edge, self.input_joint_num * channels_per_edge)
+
+        for i, pair in enumerate(self.pooling_list):
+            for j in pair:
+                for c in range(channels_per_edge):
+                    self.weight[j * channels_per_edge + c, i * channels_per_edge + c] = 1
+
+        self.weight = nn.Parameter(self.weight)
+        self.weight.requires_grad_(False)
+
+    def forward(self, input: torch.Tensor):
+        return torch.matmul(self.weight, input.unsqueeze(-1)).squeeze(-1)
+
+
+def find_neighbor_joint(parents, threshold):
+    n_joint = len(parents)
+    dist_mat = np.empty((n_joint, n_joint), dtype=np.int)
+    dist_mat[:, :] = 100000
+    for i, p in enumerate(parents):
+        dist_mat[i, i] = 0
+        if i != 0:
+            dist_mat[i, p] = dist_mat[p, i] = 1
+
+    """
+    Floyd's algorithm
+    """
+    for k in range(n_joint):
+        for i in range(n_joint):
+            for j in range(n_joint):
+                dist_mat[i, j] = min(dist_mat[i, j], dist_mat[i, k] + dist_mat[k, j])
+
+    neighbor_list = []
+    for i in range(n_joint):
+        neighbor = []
+        for j in range(n_joint):
+            if dist_mat[i, j] <= threshold:
+                neighbor.append(j)
+        neighbor_list.append(neighbor)
+
+    return neighbor_list

utils/transforms.py

@@ -0,0 +1,399 @@
+import numpy as np
+import torch
+
+
+def batch_mm(matrix, matrix_batch):
+    """
+    https://github.com/pytorch/pytorch/issues/14489#issuecomment-607730242
+    :param matrix: Sparse or dense matrix, size (m, n).
+    :param matrix_batch: Batched dense matrices, size (b, n, k).
+    :return: The batched matrix-matrix product, size (m, n) x (b, n, k) = (b, m, k).
+    """
+    batch_size = matrix_batch.shape[0]
+    # Stack the vector batch into columns. (b, n, k) -> (n, b, k) -> (n, b*k)
+    vectors = matrix_batch.transpose(0, 1).reshape(matrix.shape[1], -1)
+
+    # A matrix-matrix product is a batched matrix-vector product of the columns.
+    # And then reverse the reshaping. (m, n) x (n, b*k) = (m, b*k) -> (m, b, k) -> (b, m, k)
+    return matrix.mm(vectors).reshape(matrix.shape[0], batch_size, -1).transpose(1, 0)
+
+
+def aa2quat(rots, form='wxyz', unified_orient=True):
+    """
+    Convert angle-axis representation to wxyz quaternion and to the half plan (w >= 0)
+    @param rots: angle-axis rotations, (*, 3)
+    @param form: quaternion format, either 'wxyz' or 'xyzw'
+    @param unified_orient: Use unified orientation for quaternion (quaternion is dual cover of SO3)
+    :return:
+    """
+    angles = rots.norm(dim=-1, keepdim=True)
+    norm = angles.clone()
+    norm[norm < 1e-8] = 1
+    axis = rots / norm
+    quats = torch.empty(rots.shape[:-1] + (4,), device=rots.device, dtype=rots.dtype)
+    angles = angles * 0.5
+    if form == 'wxyz':
+        quats[..., 0] = torch.cos(angles.squeeze(-1))
+        quats[..., 1:] = torch.sin(angles) * axis
+    elif form == 'xyzw':
+        quats[..., :3] = torch.sin(angles) * axis
+        quats[..., 3] = torch.cos(angles.squeeze(-1))
+
+    if unified_orient:
+        idx = quats[..., 0] < 0
+        quats[idx, :] *= -1
+
+    return quats
+
+
+def quat2aa(quats):
+    """
+    Convert wxyz quaternions to angle-axis representation
+    :param quats:
+    :return:
+    """
+    _cos = quats[..., 0]
+    xyz = quats[..., 1:]
+    _sin = xyz.norm(dim=-1)
+    norm = _sin.clone()
+    norm[norm < 1e-7] = 1
+    axis = xyz / norm.unsqueeze(-1)
+    angle = torch.atan2(_sin, _cos) * 2
+    return axis * angle.unsqueeze(-1)
+
+
+def quat2mat(quats: torch.Tensor):
+    """
+    Convert (w, x, y, z) quaternions to 3x3 rotation matrix
+    :param quats: quaternions of shape (..., 4)
+    :return:  rotation matrices of shape (..., 3, 3)
+    """
+    qw = quats[..., 0]
+    qx = quats[..., 1]
+    qy = quats[..., 2]
+    qz = quats[..., 3]
+
+    x2 = qx + qx
+    y2 = qy + qy
+    z2 = qz + qz
+    xx = qx * x2
+    yy = qy * y2
+    wx = qw * x2
+    xy = qx * y2
+    yz = qy * z2
+    wy = qw * y2
+    xz = qx * z2
+    zz = qz * z2
+    wz = qw * z2
+
+    m = torch.empty(quats.shape[:-1] + (3, 3), device=quats.device, dtype=quats.dtype)
+    m[..., 0, 0] = 1.0 - (yy + zz)
+    m[..., 0, 1] = xy - wz
+    m[..., 0, 2] = xz + wy
+    m[..., 1, 0] = xy + wz
+    m[..., 1, 1] = 1.0 - (xx + zz)
+    m[..., 1, 2] = yz - wx
+    m[..., 2, 0] = xz - wy
+    m[..., 2, 1] = yz + wx
+    m[..., 2, 2] = 1.0 - (xx + yy)
+
+    return m
+
+
+def quat2euler(q, order='xyz', degrees=True):
+    """
+    Convert (w, x, y, z) quaternions to xyz euler angles. This is  used for bvh output.
+    """
+    q0 = q[..., 0]
+    q1 = q[..., 1]
+    q2 = q[..., 2]
+    q3 = q[..., 3]
+    es = torch.empty(q0.shape + (3,), device=q.device, dtype=q.dtype)
+
+    if order == 'xyz':
+        es[..., 2] = torch.atan2(2 * (q0 * q3 - q1 * q2), q0 * q0 + q1 * q1 - q2 * q2 - q3 * q3)
+        es[..., 1] = torch.asin((2 * (q1 * q3 + q0 * q2)).clip(-1, 1))
+        es[..., 0] = torch.atan2(2 * (q0 * q1 - q2 * q3), q0 * q0 - q1 * q1 - q2 * q2 + q3 * q3)
+    else:
+        raise NotImplementedError('Cannot convert to ordering %s' % order)
+
+    if degrees:
+        es = es * 180 / np.pi
+
+    return es
+
+
+def euler2mat(rots, order='xyz'):
+    axis = {'x': torch.tensor((1, 0, 0), device=rots.device),
+            'y': torch.tensor((0, 1, 0), device=rots.device),
+            'z': torch.tensor((0, 0, 1), device=rots.device)}
+
+    rots = rots / 180 * np.pi
+    mats = []
+    for i in range(3):
+        aa = axis[order[i]] * rots[..., i].unsqueeze(-1)
+        mats.append(aa2mat(aa))
+    return mats[0] @ (mats[1] @ mats[2])
+
+
+def aa2mat(rots):
+    """
+    Convert angle-axis representation to rotation matrix
+    :param rots: angle-axis representation
+    :return:
+    """
+    quat = aa2quat(rots)
+    mat = quat2mat(quat)
+    return mat
+
+
+def mat2quat(R) -> torch.Tensor:
+    '''
+    https://github.com/duolu/pyrotation/blob/master/pyrotation/pyrotation.py
+    Convert a rotation matrix to a unit quaternion.
+
+    This uses the Shepperd’s method for numerical stability.
+    '''
+
+    # The rotation matrix must be orthonormal
+
+    w2 = (1 + R[..., 0, 0] + R[..., 1, 1] + R[..., 2, 2])
+    x2 = (1 + R[..., 0, 0] - R[..., 1, 1] - R[..., 2, 2])
+    y2 = (1 - R[..., 0, 0] + R[..., 1, 1] - R[..., 2, 2])
+    z2 = (1 - R[..., 0, 0] - R[..., 1, 1] + R[..., 2, 2])
+
+    yz = (R[..., 1, 2] + R[..., 2, 1])
+    xz = (R[..., 2, 0] + R[..., 0, 2])
+    xy = (R[..., 0, 1] + R[..., 1, 0])
+
+    wx = (R[..., 2, 1] - R[..., 1, 2])
+    wy = (R[..., 0, 2] - R[..., 2, 0])
+    wz = (R[..., 1, 0] - R[..., 0, 1])
+
+    w = torch.empty_like(x2)
+    x = torch.empty_like(x2)
+    y = torch.empty_like(x2)
+    z = torch.empty_like(x2)
+
+    flagA = (R[..., 2, 2] < 0) * (R[..., 0, 0] > R[..., 1, 1])
+    flagB = (R[..., 2, 2] < 0) * (R[..., 0, 0] <= R[..., 1, 1])
+    flagC = (R[..., 2, 2] >= 0) * (R[..., 0, 0] < -R[..., 1, 1])
+    flagD = (R[..., 2, 2] >= 0) * (R[..., 0, 0] >= -R[..., 1, 1])
+
+    x[flagA] = torch.sqrt(x2[flagA])
+    w[flagA] = wx[flagA] / x[flagA]
+    y[flagA] = xy[flagA] / x[flagA]
+    z[flagA] = xz[flagA] / x[flagA]
+
+    y[flagB] = torch.sqrt(y2[flagB])
+    w[flagB] = wy[flagB] / y[flagB]
+    x[flagB] = xy[flagB] / y[flagB]
+    z[flagB] = yz[flagB] / y[flagB]
+
+    z[flagC] = torch.sqrt(z2[flagC])
+    w[flagC] = wz[flagC] / z[flagC]
+    x[flagC] = xz[flagC] / z[flagC]
+    y[flagC] = yz[flagC] / z[flagC]
+
+    w[flagD] = torch.sqrt(w2[flagD])
+    x[flagD] = wx[flagD] / w[flagD]
+    y[flagD] = wy[flagD] / w[flagD]
+    z[flagD] = wz[flagD] / w[flagD]
+
+    # if R[..., 2, 2] < 0:
+    #
+    #     if R[..., 0, 0] > R[..., 1, 1]:
+    #
+    #         x = torch.sqrt(x2)
+    #         w = wx / x
+    #         y = xy / x
+    #         z = xz / x
+    #
+    #     else:
+    #
+    #         y = torch.sqrt(y2)
+    #         w = wy / y
+    #         x = xy / y
+    #         z = yz / y
+    #
+    # else:
+    #
+    #     if R[..., 0, 0] < -R[..., 1, 1]:
+    #
+    #         z = torch.sqrt(z2)
+    #         w = wz / z
+    #         x = xz / z
+    #         y = yz / z
+    #
+    #     else:
+    #
+    #         w = torch.sqrt(w2)
+    #         x = wx / w
+    #         y = wy / w
+    #         z = wz / w
+
+    res = [w, x, y, z]
+    res = [z.unsqueeze(-1) for z in res]
+
+    return torch.cat(res, dim=-1) / 2
+
+
+def quat2repr6d(quat):
+    mat = quat2mat(quat)
+    res = mat[..., :2, :]
+    res = res.reshape(res.shape[:-2] + (6, ))
+    return res
+
+
+def repr6d2mat(repr):
+    x = repr[..., :3]
+    y = repr[..., 3:]
+    x = x / x.norm(dim=-1, keepdim=True)
+    z = torch.cross(x, y)
+    z = z / z.norm(dim=-1, keepdim=True)
+    y = torch.cross(z, x)
+    res = [x, y, z]
+    res = [v.unsqueeze(-2) for v in res]
+    mat = torch.cat(res, dim=-2)
+    return mat
+
+
+def repr6d2quat(repr) -> torch.Tensor:
+    x = repr[..., :3]
+    y = repr[..., 3:]
+    x = x / x.norm(dim=-1, keepdim=True)
+    z = torch.cross(x, y)
+    z = z / z.norm(dim=-1, keepdim=True)
+    y = torch.cross(z, x)
+    res = [x, y, z]
+    res = [v.unsqueeze(-2) for v in res]
+    mat = torch.cat(res, dim=-2)
+    return mat2quat(mat)
+
+
+def inv_affine(mat):
+    """
+    Calculate the inverse of any affine transformation
+    """
+    affine = torch.zeros((mat.shape[:2] + (1, 4)))
+    affine[..., 3] = 1
+    vert_mat = torch.cat((mat, affine), dim=2)
+    vert_mat_inv = torch.inverse(vert_mat)
+    return vert_mat_inv[..., :3, :]
+
+
+def inv_rigid_affine(mat):
+    """
+    Calculate the inverse of a rigid affine transformation
+    """
+    res = mat.clone()
+    res[..., :3] = mat[..., :3].transpose(-2, -1)
+    res[..., 3] = -torch.matmul(res[..., :3], mat[..., 3].unsqueeze(-1)).squeeze(-1)
+    return res
+
+
+def generate_pose(batch_size, device, uniform=False, factor=1, root_rot=False, n_bone=None, ee=None):
+    if n_bone is None: n_bone = 24
+    if ee is not None:
+        if root_rot:
+            ee.append(0)
+        n_bone_ = n_bone
+        n_bone = len(ee)
+    axis = torch.randn((batch_size, n_bone, 3), device=device)
+    axis /= axis.norm(dim=-1, keepdim=True)
+    if uniform:
+        angle = torch.rand((batch_size, n_bone, 1), device=device) * np.pi
+    else:
+        angle = torch.randn((batch_size, n_bone, 1), device=device) * np.pi / 6 * factor
+        angle.clamp(-np.pi, np.pi)
+    poses = axis * angle
+    if ee is not None:
+        res = torch.zeros((batch_size, n_bone_, 3), device=device)
+        for i, id in enumerate(ee):
+            res[:, id] = poses[:, i]
+        poses = res
+    poses = poses.reshape(batch_size, -1)
+    if not root_rot:
+        poses[..., :3] = 0
+    return poses
+
+
+def slerp(l, r, t, unit=True):
+    """
+    :param l: shape = (*, n)
+    :param r: shape = (*, n)
+    :param t: shape = (*)
+    :param unit: If l and h are unit vectors
+    :return:
+    """
+    eps = 1e-8
+    if not unit:
+        l_n = l / torch.norm(l, dim=-1, keepdim=True)
+        r_n = r / torch.norm(r, dim=-1, keepdim=True)
+    else:
+        l_n = l
+        r_n = r
+    omega = torch.acos((l_n * r_n).sum(dim=-1).clamp(-1, 1))
+    dom = torch.sin(omega)
+
+    flag = dom < eps
+
+    res = torch.empty_like(l_n)
+    t_t = t[flag].unsqueeze(-1)
+    res[flag] = (1 - t_t) * l_n[flag] + t_t * r_n[flag]
+
+    flag = ~ flag
+
+    t_t = t[flag]
+    d_t = dom[flag]
+    va = torch.sin((1 - t_t) * omega[flag]) / d_t
+    vb = torch.sin(t_t * omega[flag]) / d_t
+    res[flag] = (va.unsqueeze(-1) * l_n[flag] + vb.unsqueeze(-1) * r_n[flag])
+    return res
+
+
+def slerp_quat(l, r, t):
+    """
+    slerp for unit quaternions
+    :param l: (*, 4) unit quaternion
+    :param r: (*, 4) unit quaternion
+    :param t: (*) scalar between 0 and 1
+    """
+    t = t.expand(l.shape[:-1])
+    flag = (l * r).sum(dim=-1) >= 0
+    res = torch.empty_like(l)
+    res[flag] = slerp(l[flag], r[flag], t[flag])
+    flag = ~ flag
+    res[flag] = slerp(-l[flag], r[flag], t[flag])
+    return res
+
+
+# def slerp_6d(l, r, t):
+#     l_q = repr6d2quat(l)
+#     r_q = repr6d2quat(r)
+#     res_q = slerp_quat(l_q, r_q, t)
+#     return quat2repr6d(res_q)
+
+
+def interpolate_6d(input, size):
+    """
+    :param input: (batch_size, n_channels, length)
+    :param size: required output size for temporal axis
+    :return:
+    """
+    batch = input.shape[0]
+    length = input.shape[-1]
+    input = input.reshape((batch, -1, 6, length))
+    input = input.permute(0, 1, 3, 2)     # (batch_size, n_joint, length, 6)
+    input_q = repr6d2quat(input)
+    idx = torch.tensor(list(range(size)), device=input_q.device, dtype=torch.float) / size * (length - 1)
+    idx_l = torch.floor(idx)
+    t = idx - idx_l
+    idx_l = idx_l.long()
+    idx_r = idx_l + 1
+    t = t.reshape((1, 1, -1))
+    res_q = slerp_quat(input_q[..., idx_l, :], input_q[..., idx_r, :], t)
+    res = quat2repr6d(res_q)  # shape = (batch_size, n_joint, t, 6)
+    res = res.permute(0, 1, 3, 2)
+    res = res.reshape((batch, -1, size))
+    return res