jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def gaussian_curvature(grad: torch.Tensor, hess: torch.Tensor) -> torch.Tensor:
    """Calculates the Gaussian curvature of a implicit surface.
    See https://en.wikipedia.org/wiki/Gaussian_curvature#Alternative_formulas
    for details.

    Parameters
    ----------
    grad: torch.Tensor
        Gradient of the surface, shaped [B, N, 3].

    hess: torch.Tensor
        Hessian of the surface, shaped [B, N, 1, 1, 3, 3].

    Returns
    -------
    Kg: torch.Tensor
        The gaussian curvatures.
    """
    # Append gradients to the last columns of the hessians.
    grad5d = torch.unsqueeze(grad, 2)
    grad5d = torch.unsqueeze(grad5d, -1)
    F = torch.cat((hess, grad5d), -1)

    # Append gradients (with and additional 0 at the last coord) to the last
    # lines of the hessians.
    hess_size = hess.size()
    zeros_size = list(itertools.chain.from_iterable((hess_size[:3], [1, 1])))
    zeros = torch.zeros(zeros_size).to(grad.device)
    grad5d = torch.unsqueeze(grad, 2)
    grad5d = torch.unsqueeze(grad5d, -2)
    grad5d = torch.cat((grad5d, zeros), -1)

    F = torch.cat((F, grad5d), -2)
    grad_norm = torch.norm(grad, dim=-1)

    Kg = -torch.det(F)[-1].squeeze(-1) / (grad_norm[0]**4)
    return Kg
//  [python]
def gaussian_curvature(grad: torch.Tensor, hess: torch.Tensor) -> torch.Tensor:
    """Calculates the Gaussian curvature of a implicit surface.
    See https://en.wikipedia.org/wiki/Gaussian_curvature#Alternative_formulas
    for details.

    Parameters
    ----------
    grad: torch.Tensor
        Gradient of the surface, shaped [B, N, 3].

    hess: torch.Tensor
        Hessian of the surface, shaped [B, N, 1, 1, 3, 3].

    Returns
    -------
    Kg: torch.Tensor
        The gaussian curvatures.
    """
    # Append gradients to the last columns of the hessians.
    grad5d = torch.unsqueeze(grad, 2)
    grad5d = torch.unsqueeze(grad5d, -1)
    F = torch.cat((hess, grad5d), -1)

    # Append gradients (with and additional 0 at the last coord) to the last
    # lines of the hessians.
    hess_size = hess.size()
    zeros_size = list(itertools.chain.from_iterable((hess_size[:3], [1, 1])))
    zeros = torch.zeros(zeros_size).to(grad.device)
    grad5d = torch.unsqueeze(grad, 2)
    grad5d = torch.unsqueeze(grad5d, -2)
    grad5d = torch.cat((grad5d, zeros), -1)

    F = torch.cat((F, grad5d), -2)
    grad_norm = torch.linalg.norm(grad, dim=-1)

    Kg = -torch.det(F)[-1].squeeze(-1) / (grad_norm[0]**4)
    return Kg
//  [python]
def closedPathAsBezierSpline(Y : torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
    if not torch.all(Y[0]==Y[-1]):
        Yaug : torch.Tensor = torch.cat([Y, Y[0].unsqueeze(0)], dim=0)
    else:
        Yaug : torch.Tensor = Y
    euclidean_distances : torch.Tensor = torch.zeros_like(Yaug[:,0])
    delta_euclidean_distances = torch.norm( Yaug[1:]-Yaug[:-1] , dim=1 )
    euclidean_distances[1:] = torch.cumsum( delta_euclidean_distances, 0 )
    euclidean_spline = compositeBezierSpline(euclidean_distances,Yaug,boundary_conditions="periodic")
    arclengths = torch.zeros_like(Yaug[:,0])
    distances = bezierArcLength(euclidean_spline)
    arclengths[1:] = torch.cumsum(distances, 0)
    return arclengths, compositeBezierSpline(arclengths,Yaug,boundary_conditions="periodic")
//  [python]
def closedPathAsBezierSpline(Y : torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
    if not torch.all(Y[0]==Y[-1]):
        Yaug : torch.Tensor = torch.cat([Y, Y[0].unsqueeze(0)], dim=0)
    else:
        Yaug : torch.Tensor = Y
    euclidean_distances : torch.Tensor = torch.zeros_like(Yaug[:,0])
    delta_euclidean_distances = torch.linalg.norm( Yaug[1:]-Yaug[:-1] , dim=1 )
    euclidean_distances[1:] = torch.cumsum( delta_euclidean_distances, 0 )
    euclidean_spline = compositeBezierSpline(euclidean_distances,Yaug,boundary_conditions="periodic")
    arclengths = torch.zeros_like(Yaug[:,0])
    distances = bezierArcLength(euclidean_spline)
    arclengths[1:] = torch.cumsum(distances, 0)
    return arclengths, compositeBezierSpline(arclengths,Yaug,boundary_conditions="periodic")
//  [python]
def test_lennard_jones_batched_jac(self, device, jac):
        sigma = 0.5
        epsilon = 4.0

        jac = getattr(functorch, jac)

        def lennard_jones(r):
            return epsilon * ((sigma / r) ** 12 - (sigma / r) ** 6)

        def lennard_jones_force(r):
            """Get magnitude of LJ force"""
            return -epsilon * (
                (-12 * sigma**12 / r**13) + (6 * sigma**6 / r**7)
            )

        r = torch.linspace(0.5, 2 * sigma, steps=100, requires_grad=True, device=device)
        drs = torch.outer(r, torch.tensor([1.0, 0, 0], device=device))
        norms = torch.norm(drs, dim=1).reshape(-1, 1)
        training_energies = torch.stack(list(map(lennard_jones, norms))).reshape(-1, 1)
        training_forces = torch.stack(
            [force * dr for force, dr in zip(map(lennard_jones_force, norms), drs)]
        )

        model = nn.Sequential(
            nn.Linear(1, 16),
            nn.Tanh(),
            nn.Linear(16, 16),
            nn.Tanh(),
            nn.Linear(16, 16),
            nn.Tanh(),
            nn.Linear(16, 16),
            nn.Tanh(),
            nn.Linear(16, 1),
        ).to(device)

        def make_prediction(model, drs, use_functorch):
            norms = torch.norm(drs, dim=1).reshape(-1, 1)
            energies = model(norms)

            if use_functorch:
                network_derivs = vmap(jac(model))(norms).squeeze(-1)
                forces = -network_derivs * drs / norms
            else:
                forces = []
                for r, dr in zip(norms, drs):
                    network_deriv = torch.autograd.functional.jacobian(
                        model, r, create_graph=True
                    )
                    force = -network_deriv * dr / r
                    forces.append(force)
                forces = torch.cat(forces)
            return energies, forces

        def loss_fn(energies, forces, predicted_energies, predicted_forces):
            return (
                F.mse_loss(energies, predicted_energies)
                + 0.01 * F.mse_loss(forces, predicted_forces) / 3
            )

        energies, forces = make_prediction(model, drs, use_functorch=True)
        loss = loss_fn(training_energies, training_forces, energies, forces)
        result = torch.autograd.grad(loss, model.parameters())

        energies, forces = make_prediction(model, drs, use_functorch=False)
        loss = loss_fn(training_energies, training_forces, energies, forces)
        expected = torch.autograd.grad(loss, model.parameters())

        self.assertEqual(result, expected)
//  [python]
def test_lennard_jones_batched_jac(self, device, jac):
        sigma = 0.5
        epsilon = 4.0

        jac = getattr(functorch, jac)

        def lennard_jones(r):
            return epsilon * ((sigma / r) ** 12 - (sigma / r) ** 6)

        def lennard_jones_force(r):
            """Get magnitude of LJ force"""
            return -epsilon * (
                (-12 * sigma**12 / r**13) + (6 * sigma**6 / r**7)
            )

        r = torch.linspace(0.5, 2 * sigma, steps=100, requires_grad=True, device=device)
        drs = torch.outer(r, torch.tensor([1.0, 0, 0], device=device))
        norms = torch.linalg.norm(drs, dim=1).reshape(-1, 1)
        training_energies = torch.stack(list(map(lennard_jones, norms))).reshape(-1, 1)
        training_forces = torch.stack(
            [force * dr for force, dr in zip(map(lennard_jones_force, norms), drs)]
        )

        model = nn.Sequential(
            nn.Linear(1, 16),
            nn.Tanh(),
            nn.Linear(16, 16),
            nn.Tanh(),
            nn.Linear(16, 16),
            nn.Tanh(),
            nn.Linear(16, 16),
            nn.Tanh(),
            nn.Linear(16, 1),
        ).to(device)

        def make_prediction(model, drs, use_functorch):
            norms = torch.linalg.norm(drs, dim=1).reshape(-1, 1)
            energies = model(norms)

            if use_functorch:
                network_derivs = vmap(jac(model))(norms).squeeze(-1)
                forces = -network_derivs * drs / norms
            else:
                forces = []
                for r, dr in zip(norms, drs):
                    network_deriv = torch.autograd.functional.jacobian(
                        model, r, create_graph=True
                    )
                    force = -network_deriv * dr / r
                    forces.append(force)
                forces = torch.cat(forces)
            return energies, forces

        def loss_fn(energies, forces, predicted_energies, predicted_forces):
            return (
                F.mse_loss(energies, predicted_energies)
                + 0.01 * F.mse_loss(forces, predicted_forces) / 3
            )

        energies, forces = make_prediction(model, drs, use_functorch=True)
        loss = loss_fn(training_energies, training_forces, energies, forces)
        result = torch.autograd.grad(loss, model.parameters())

        energies, forces = make_prediction(model, drs, use_functorch=False)
        loss = loss_fn(training_energies, training_forces, energies, forces)
        expected = torch.autograd.grad(loss, model.parameters())

        self.assertEqual(result, expected)
//  [python]
def exp2quat(e):
    """
    Convert axis-angle rotations (aka exponential maps) to quaternions.
    Stable formula from "Practical Parameterization of Rotations Using the Exponential Map".
    Expects a tensor of shape (*, 3), where * denotes any number of dimensions.
    Returns a tensor of shape (*, 4).
    """
    assert e.shape[-1] == 3 
    original_shape = list(e.shape)
    original_shape[-1] = 4
    e = e.reshape(-1, 3) 
    theta = torch.norm(e,dim=1).view(-1, 1) 
    w = torch.cos(0.5 * theta).view(-1, 1) 
    xyz = 0.5 * pysinc(0.5 * theta / math.pi) * e
    return torch.cat((w, xyz), 1).view(original_shape)
//  [python]
def exp2quat(e):
    """
    Convert axis-angle rotations (aka exponential maps) to quaternions.
    Stable formula from "Practical Parameterization of Rotations Using the Exponential Map".
    Expects a tensor of shape (*, 3), where * denotes any number of dimensions.
    Returns a tensor of shape (*, 4).
    """
    assert e.shape[-1] == 3 
    original_shape = list(e.shape)
    original_shape[-1] = 4
    e = e.reshape(-1, 3) 
    theta = torch.linalg.norm(e,dim=1).view(-1, 1) 
    w = torch.cos(0.5 * theta).view(-1, 1) 
    xyz = 0.5 * pysinc(0.5 * theta / math.pi) * e
    return torch.cat((w, xyz), 1).view(original_shape)
//  [python]
def valid_flags_by_vnect_accuracy(th_error,th_count,data,gt):
    N, T, J, _ = gt.shape
    gt = gt.view(N, T, -1)
    data = data.view(N, T, -1)
    residual = gt - data
    norms = torch.norm(residual, dim=2) 
    labels = torch.zeros(norms.shape)
    higher_id = torch.nonzero(norms > th_error) 
    labels[higher_id[:, 0], higher_id[:, 1]] = 1

    labels_sequence_wise = torch.sum(labels, 1)
    higher_id_sequence_wise = torch.nonzero(labels_sequence_wise > th_count)
    valid_labels = torch.ones(labels_sequence_wise.shape)
    valid_labels[higher_id_sequence_wise[:, 0]] = 0
    return valid_labels
//  [python]
def valid_flags_by_vnect_accuracy(th_error,th_count,data,gt):
    N, T, J, _ = gt.shape
    gt = gt.view(N, T, -1)
    data = data.view(N, T, -1)
    residual = gt - data
    norms = torch.linalg.norm(residual, dim=2) 
    labels = torch.zeros(norms.shape)
    higher_id = torch.nonzero(norms > th_error) 
    labels[higher_id[:, 0], higher_id[:, 1]] = 1

    labels_sequence_wise = torch.sum(labels, 1)
    higher_id_sequence_wise = torch.nonzero(labels_sequence_wise > th_count)
    valid_labels = torch.ones(labels_sequence_wise.shape)
    valid_labels[higher_id_sequence_wise[:, 0]] = 0
    return valid_labels
//  [python]
def add_batch(self, inp, out):
        if len(inp.shape) == 2:
            inp = inp.unsqueeze(0)
        tmp = inp.shape[0]
        if isinstance(self.layer, nn.Linear) or isinstance(
            self.layer, transformers.Conv1D
        ):
            if len(inp.shape) == 3:
                inp = inp.reshape((-1, inp.shape[-1]))
            inp = inp.t()
        self.H *= self.nsamples / (self.nsamples + tmp)

        self.scaler_row *= self.nsamples / (self.nsamples + tmp)

        self.nsamples += tmp
        inp = math.sqrt(2 / self.nsamples) * inp.float()
        self.H += inp.matmul(inp.t())
        self.scaler_row += torch.norm(inp, p=2, dim=1) ** 2 / self.nsamples
//  [python]
def add_batch(self, inp, out):
        if len(inp.shape) == 2:
            inp = inp.unsqueeze(0)
        tmp = inp.shape[0]
        if isinstance(self.layer, nn.Linear) or isinstance(
            self.layer, transformers.Conv1D
        ):
            if len(inp.shape) == 3:
                inp = inp.reshape((-1, inp.shape[-1]))
            inp = inp.t()
        self.H *= self.nsamples / (self.nsamples + tmp)

        self.scaler_row *= self.nsamples / (self.nsamples + tmp)

        self.nsamples += tmp
        inp = math.sqrt(2 / self.nsamples) * inp.float()
        self.H += inp.matmul(inp.t())
        self.scaler_row += torch.linalg.norm(inp, p=2, dim=1) ** 2 / self.nsamples
//  [python]
def bezierArcLength(control_points : torch.Tensor, quadrature_order = 7, num_segments = 4, sum=True):
    
    batchdim = control_points.shape[0]
    kbezier = control_points.shape[-2] - 1
    dim = control_points.shape[-1]
    
    dtype = control_points.dtype
    device = control_points.device

    ds = 1.0/num_segments
    
    a = torch.linspace(0.0, 1.0 - ds, steps=num_segments, dtype=dtype, device=device)
    b = a + ds

    roots_np, weights_np = roots_legendre(quadrature_order)
    roots = torch.as_tensor(roots_np, dtype=dtype, device=device)
    weights = torch.as_tensor(weights_np, dtype=dtype, device=device)

    s_shifted = 0.5*ds*roots.unsqueeze(0).expand(num_segments, quadrature_order) + 0.5*(a+b)[:,None]

    bezierMderiv = bezierM(s_shifted, kbezier-1)
    bezierMderiv_exp = bezierMderiv.unsqueeze(0).expand(batchdim, num_segments, quadrature_order, kbezier)

    control_points_deriv = kbezier*(control_points[:,1:] - control_points[:,:-1])
    control_points_deriv_exp = control_points_deriv.unsqueeze(1).expand(batchdim, num_segments, kbezier, dim)

    quadrature_node_values = torch.matmul(bezierMderiv_exp, control_points_deriv_exp)
    quadrature_node_norms = torch.norm(quadrature_node_values, p=2.0, dim=-1)

    segment_sums = 0.5*ds*torch.sum(quadrature_node_norms*weights[None,None], dim=-1)

    if sum:
        return torch.sum(segment_sums, dim=1)
    
    return segment_sums
//  [python]
def bezierArcLength(control_points : torch.Tensor, quadrature_order = 7, num_segments = 4, sum=True):
    
    batchdim = control_points.shape[0]
    kbezier = control_points.shape[-2] - 1
    dim = control_points.shape[-1]
    
    dtype = control_points.dtype
    device = control_points.device

    ds = 1.0/num_segments
    
    a = torch.linspace(0.0, 1.0 - ds, steps=num_segments, dtype=dtype, device=device)
    b = a + ds

    roots_np, weights_np = roots_legendre(quadrature_order)
    roots = torch.as_tensor(roots_np, dtype=dtype, device=device)
    weights = torch.as_tensor(weights_np, dtype=dtype, device=device)

    s_shifted = 0.5*ds*roots.unsqueeze(0).expand(num_segments, quadrature_order) + 0.5*(a+b)[:,None]

    bezierMderiv = bezierM(s_shifted, kbezier-1)
    bezierMderiv_exp = bezierMderiv.unsqueeze(0).expand(batchdim, num_segments, quadrature_order, kbezier)

    control_points_deriv = kbezier*(control_points[:,1:] - control_points[:,:-1])
    control_points_deriv_exp = control_points_deriv.unsqueeze(1).expand(batchdim, num_segments, kbezier, dim)

    quadrature_node_values = torch.matmul(bezierMderiv_exp, control_points_deriv_exp)
    quadrature_node_norms = torch.linalg.norm(quadrature_node_values, p=2.0, dim=-1)

    segment_sums = 0.5*ds*torch.sum(quadrature_node_norms*weights[None,None], dim=-1)

    if sum:
        return torch.sum(segment_sums, dim=1)
    
    return segment_sums
//  [python]
def qinv(q: torch.Tensor):
    """
    Invert quaternions
    Expect a tensor of shape (*, 4)
    Returns a tensor of shape (*, 4)
    """
    ori_shape = q.shape
    q = q.contiguous().view(-1, 4)
    q_norm = torch.norm(q, p=2, dim=1)
    q_inv = q.clone().contiguous().view(-1, 4)
    q_inv[:, 1:] *= -1
    q_inv /= q_norm.unsqueeze(1)
    return q_inv.reshape(*ori_shape)
//  [python]
def qinv(q: torch.Tensor):
    """
    Invert quaternions
    Expect a tensor of shape (*, 4)
    Returns a tensor of shape (*, 4)
    """
    ori_shape = q.shape
    q = q.contiguous().view(-1, 4)
    q_norm = torch.linalg.norm(q, ord=2, dim=1)
    q_inv = q.clone().contiguous().view(-1, 4)
    q_inv[:, 1:] *= -1
    q_inv /= q_norm.unsqueeze(1)
    return q_inv.reshape(*ori_shape)
//  [python]
def ISTA(self, D, x, h, alpha, lamb):
        critical_value = (torch.eig(torch.mm(D.t().to(x.device), D.to(x.device)))[0][:, 0]).max()

        if alpha * critical_value > 1:
            raise ValueError('Your alpha is too big for ISTA to converge.')
        converged = False
        iter = 0
        while not converged:
            h_prime = h
            h = h - torch.mm(alpha * D.t().to(x.device), torch.mm(D.to(x.device), h.to(x.device)) - x)
            h = self.shrink(h, alpha * lamb)
            # h = torch.clamp(h, min=0) ## every iter clamp negative values to 0
            if torch.norm(h - h_prime) < 0.0008:
                converged = True
            iter += 1
            if iter>2000: #2000
                break
        return h
//  [python]
def ISTA(self, D, x, h, alpha, lamb):
        critical_value = (torch.eig(torch.mm(D.t().to(x.device), D.to(x.device)))[0][:, 0]).max()

        if alpha * critical_value > 1:
            raise ValueError('Your alpha is too big for ISTA to converge.')
        converged = False
        iter = 0
        while not converged:
            h_prime = h
            h = h - torch.mm(alpha * D.t().to(x.device), torch.mm(D.to(x.device), h.to(x.device)) - x)
            h = self.shrink(h, alpha * lamb)
            # h = torch.clamp(h, min=0) ## every iter clamp negative values to 0
            if torch.linalg.norm(h - h_prime) < 0.0008:
                converged = True
            iter += 1
            if iter>2000: #2000
                break
        return h
//  [python]
def mpjpe_cal(predicted, target):
    assert predicted.shape == target.shape
    return torch.mean(torch.norm(predicted - target, dim=len(target.shape) - 1))
//  [python]
def mpjpe_cal(predicted, target):
    assert predicted.shape == target.shape
    return torch.mean(torch.linalg.norm(predicted - target, dim=len(target.shape) - 1))
//  [python]
def rotation_from_to(v_from: torch.Tensor, v_to: torch.Tensor):
    """
    Calculate the shortest rotation from two vectors
    Argument:
     -- Both v_from, v_to is of shape (*, 3)
    """
    assert v_from.shape == v_to.shape
    assert v_from.shape[-1] == 3 and v_to.shape[-1] == 3

    ori_shape = v_from.shape
    out_shape = list(ori_shape)[:-1] + [4]
    v_from = v_from.view(-1, 3)
    v_to = v_to.view(-1, 3)
    dot = torch.sum(v_from * v_to, dim=1)
    xyz = torch.cross(v_from, v_to)
    w = torch.norm(v_from, dim=1, p=2) * torch.norm(v_to, dim=1, p=2) + dot
    w = w.unsqueeze(1)
    rotations = torch.cat([w, xyz], dim=1)
    rotations = F.normalize(rotations, p=2, dim=1)
    return rotations.view(out_shape)
//  [python]
def rotation_from_to(v_from: torch.Tensor, v_to: torch.Tensor):
    """
    Calculate the shortest rotation from two vectors
    Argument:
     -- Both v_from, v_to is of shape (*, 3)
    """
    assert v_from.shape == v_to.shape
    assert v_from.shape[-1] == 3 and v_to.shape[-1] == 3

    ori_shape = v_from.shape
    out_shape = list(ori_shape)[:-1] + [4]
    v_from = v_from.view(-1, 3)
    v_to = v_to.view(-1, 3)
    dot = torch.sum(v_from * v_to, dim=1)
    xyz = torch.cross(v_from, v_to)
    w = torch.linalg.norm(v_from, dim=1, ord=2) * torch.linalg.norm(v_to, dim=1, ord=2) + dot
    w = w.unsqueeze(1)
    rotations = torch.cat([w, xyz], dim=1)
    rotations = F.normalize(rotations, p=2, dim=1)
    return rotations.view(out_shape)
//  [python]
def reproject_error(pt3d, pt2d, pose_pred, k_matrix):
    k_matrix = torch.from_numpy(k_matrix).double()
    f, ppx, ppy = k_matrix[0, 0], k_matrix[0, 2], k_matrix[1, 2]
    pt3d = torch.from_numpy(pt3d).permute(1, 0).double()
    pt2d = torch.from_numpy(pt2d).double()
    pose_pred = torch.from_numpy(pose_pred)
    R = pose_pred[:3, :3]
    t = pose_pred[:3, 3]
    objMat = torch.matmul(R, pt3d).permute(1, 0) + t
    obj_pt = torch.zeros_like(objMat[..., :2]).double()
    obj_pt[..., 0] = objMat[..., 0] * f / objMat[..., 2] + ppx
    obj_pt[..., 1] = objMat[..., 1] * f / objMat[..., 2] + ppy
    error = torch.norm(pt2d - obj_pt, p=2, dim=1).clamp(0, 100).mean()
    return error.float()
//  [python]
def reproject_error(pt3d, pt2d, pose_pred, k_matrix):
    k_matrix = torch.from_numpy(k_matrix).double()
    f, ppx, ppy = k_matrix[0, 0], k_matrix[0, 2], k_matrix[1, 2]
    pt3d = torch.from_numpy(pt3d).permute(1, 0).double()
    pt2d = torch.from_numpy(pt2d).double()
    pose_pred = torch.from_numpy(pose_pred)
    R = pose_pred[:3, :3]
    t = pose_pred[:3, 3]
    objMat = torch.matmul(R, pt3d).permute(1, 0) + t
    obj_pt = torch.zeros_like(objMat[..., :2]).double()
    obj_pt[..., 0] = objMat[..., 0] * f / objMat[..., 2] + ppx
    obj_pt[..., 1] = objMat[..., 1] * f / objMat[..., 2] + ppy
    error = torch.linalg.norm(pt2d - obj_pt, ord=2, dim=1).clamp(0, 100).mean()
    return error.float()
//  [python]
def mpjpe_by_action_p1(predicted, target, action, action_error_sum):
    assert predicted.shape == target.shape
    num = predicted.size(0)
    dist = torch.mean(torch.norm(predicted - target, dim=len(target.shape) - 1), dim=len(target.shape) - 2)

    if len(set(list(action))) == 1:
        end_index = action[0].find(' ')
        if end_index != -1:
            action_name = action[0][:end_index]
        else:
            action_name = action[0]

        action_error_sum[action_name]['p1'].update(torch.mean(dist).item()*num, num)
    else:
        for i in range(num):
            end_index = action[i].find(' ')
            if end_index != -1:
                action_name = action[i][:end_index]
            else:
                action_name = action[i]

            action_error_sum[action_name]['p1'].update(dist[i].item(), 1)
            
    return action_error_sum
//  [python]
def mpjpe_by_action_p1(predicted, target, action, action_error_sum):
    assert predicted.shape == target.shape
    num = predicted.size(0)
    dist = torch.mean(torch.linalg.norm(predicted - target, dim=len(target.shape) - 1), dim=len(target.shape) - 2)

    if len(set(list(action))) == 1:
        end_index = action[0].find(' ')
        if end_index != -1:
            action_name = action[0][:end_index]
        else:
            action_name = action[0]

        action_error_sum[action_name]['p1'].update(torch.mean(dist).item()*num, num)
    else:
        for i in range(num):
            end_index = action[i].find(' ')
            if end_index != -1:
                action_name = action[i][:end_index]
            else:
                action_name = action[i]

            action_error_sum[action_name]['p1'].update(dist[i].item(), 1)
            
    return action_error_sum
//  [python]
def train_lqr_value_approximator(lyapunov_relu, V_lambda, R, x_equilibrium,
                                 x_lo, x_up, num_samples, lqr_S: torch.Tensor):
    """
    We train both lyapunov_relu and R such that ϕ(x) − ϕ(x*) + λ|R(x−x*)|₁
    approximates the lqr cost-to-go.
    """
    x_samples = utils.uniform_sample_in_box(x_lo, x_up, num_samples)
    V_samples = torch.sum((x_samples.T - x_equilibrium.reshape(
        (6, 1))) * (lqr_S @ (x_samples.T - x_equilibrium.reshape((6, 1)))),
                          dim=0).reshape((-1, 1))
    state_value_dataset = torch.utils.data.TensorDataset(x_samples, V_samples)
    R.requires_grad_(True)

    def compute_v(model, x):
        return model(x) - model(x_equilibrium) + V_lambda * torch.norm(
            R @ (x - x_equilibrium.reshape((1, 6))).T, p=1, dim=0).reshape(
                (-1, 1))

    utils.train_approximator(state_value_dataset,
                             lyapunov_relu,
                             compute_v,
                             batch_size=50,
                             num_epochs=200,
                             lr=0.001,
                             additional_variable=[R])
    R.requires_grad_(False)
//  [python]
def train_lqr_value_approximator(lyapunov_relu, V_lambda, R, x_equilibrium,
                                 x_lo, x_up, num_samples, lqr_S: torch.Tensor):
    """
    We train both lyapunov_relu and R such that ϕ(x) − ϕ(x*) + λ|R(x−x*)|₁
    approximates the lqr cost-to-go.
    """
    x_samples = utils.uniform_sample_in_box(x_lo, x_up, num_samples)
    V_samples = torch.sum((x_samples.T - x_equilibrium.reshape(
        (6, 1))) * (lqr_S @ (x_samples.T - x_equilibrium.reshape((6, 1)))),
                          dim=0).reshape((-1, 1))
    state_value_dataset = torch.utils.data.TensorDataset(x_samples, V_samples)
    R.requires_grad_(True)

    def compute_v(model, x):
        return model(x) - model(x_equilibrium) + V_lambda * torch.linalg.norm(
            R @ (x - x_equilibrium.reshape((1, 6))).T, ord=1, dim=0).reshape(
                (-1, 1))

    utils.train_approximator(state_value_dataset,
                             lyapunov_relu,
                             compute_v,
                             batch_size=50,
                             num_epochs=200,
                             lr=0.001,
                             additional_variable=[R])
    R.requires_grad_(False)
//  [python]
def forward(self, hidden_states: torch.FloatTensor) -> torch.FloatTensor:
        # Compute and normalize global spatial feature maps
        global_features = torch.norm(hidden_states, p=2, dim=(1, 2), keepdim=True)
        norm_features = global_features / (global_features.mean(dim=-1, keepdim=True) + 1e-6)
        hidden_states = self.weight * (hidden_states * norm_features) + self.bias + hidden_states

        return hidden_states
//  [python]
def forward(self, hidden_states: torch.FloatTensor) -> torch.FloatTensor:
        # Compute and normalize global spatial feature maps
        global_features = torch.linalg.norm(hidden_states, ord=2, dim=(1, 2), keepdim=True)
        norm_features = global_features / (global_features.mean(dim=-1, keepdim=True) + 1e-6)
        hidden_states = self.weight * (hidden_states * norm_features) + self.bias + hidden_states

        return hidden_states
//  [python]
def sample_points_within_boxes(boxes: torch.Tensor, num_samples: int = 10):
    """
    Sample points within a box defined by its 8 coordinates
    Args:
        box: tensor of shape (N, 8, 3) of the box coordinates
        num_samples: int defining the number of samples per bounding box
    Returns:
        points: (N, num_samples, 3) of points inside the box
    """
    assert boxes.shape[1] == 8 and boxes.shape[2] == 3
    N = boxes.shape[0]

    U = boxes[:, 1:2] - boxes[:, 0:1] # (N, 1, 3)
    V = boxes[:, 2:3] - boxes[:, 1:2] # (N, 1, 3)
    W = boxes[:, 4:5] - boxes[:, 0:1] # (N, 1, 3)

    UVW = torch.cat([U, V, W], dim=1) # (N, 3, 3)

    # sample points in unit cube
    whl = torch.norm(UVW, dim=2, keepdim=True)
    points = torch.rand((N, num_samples, 3), device=boxes.device) * whl[:, None, ..., 0]
    ones_vec = torch.ones(N, num_samples, 1).to(boxes.device)
    points = torch.cat([points, ones_vec], axis=-1)

    UVW = UVW / (whl + 1e-6) # (N, 3, 3)
    transformation_matrices = torch.zeros((N, 4, 4)).to(boxes.device)
    transformation_matrices[:, -1, -1] = 1.0
    transformation_matrices[:, :-1, -1] = boxes[:, 0]
    transformation_matrices[:, :-1, :-1] = UVW.permute(0, 2, 1)

    points = points @ transformation_matrices.permute(0, 2, 1)
    return points[..., :-1]
//  [python]
def sample_points_within_boxes(boxes: torch.Tensor, num_samples: int = 10):
    """
    Sample points within a box defined by its 8 coordinates
    Args:
        box: tensor of shape (N, 8, 3) of the box coordinates
        num_samples: int defining the number of samples per bounding box
    Returns:
        points: (N, num_samples, 3) of points inside the box
    """
    assert boxes.shape[1] == 8 and boxes.shape[2] == 3
    N = boxes.shape[0]

    U = boxes[:, 1:2] - boxes[:, 0:1] # (N, 1, 3)
    V = boxes[:, 2:3] - boxes[:, 1:2] # (N, 1, 3)
    W = boxes[:, 4:5] - boxes[:, 0:1] # (N, 1, 3)

    UVW = torch.cat([U, V, W], dim=1) # (N, 3, 3)

    # sample points in unit cube
    whl = torch.linalg.norm(UVW, dim=2, keepdim=True)
    points = torch.rand((N, num_samples, 3), device=boxes.device) * whl[:, None, ..., 0]
    ones_vec = torch.ones(N, num_samples, 1).to(boxes.device)
    points = torch.cat([points, ones_vec], axis=-1)

    UVW = UVW / (whl + 1e-6) # (N, 3, 3)
    transformation_matrices = torch.zeros((N, 4, 4)).to(boxes.device)
    transformation_matrices[:, -1, -1] = 1.0
    transformation_matrices[:, :-1, -1] = boxes[:, 0]
    transformation_matrices[:, :-1, :-1] = UVW.permute(0, 2, 1)

    points = points @ transformation_matrices.permute(0, 2, 1)
    return points[..., :-1]
//  [python]
def estimate_initial_trans(body_pose, joints3d_op, joints2d_op, focal):
    """
    use focal length and bone lengths to approximate distance from camera
    (based on PROX https://github.com/mohamedhassanmus/prox/blob/master/prox/fitting.py)
    :param body_pose
    :param joints3d_op OP keypoints in 3d
    :param joints2d_op OP keypoints in 2d
    """
    B, T, N, _ = joints2d_op.shape
    joints2d_obs = joints2d_op[:, :, :, :2]  # (B, T, N, 3)
    joints2d_conf = joints2d_op[:, :, :, 2]

    # find least-occluded 2d frame
    num_2d_vis = torch.sum(joints2d_conf > 0.0, dim=2)
    best_2d_idx = torch.max(num_2d_vis, dim=1)[1]

    # calculate bone lengths and confidence in each bone length
    bone3d = []
    bone2d = []
    conf2d = []
    for pair in OP_EDGE_LIST:
        diff3d = torch.norm(
            joints3d_op[:, 0, pair[0], :] - joints3d_op[:, 0, pair[1], :],
            dim=1,
        )  # does not change over time
        diff2d = torch.norm(
            joints2d_obs[:, :, pair[0], :] - joints2d_obs[:, :, pair[1], :],
            dim=2,
        )
        minconf2d = torch.min(
            joints2d_conf[:, :, pair[0]], joints2d_conf[:, :, pair[1]]
        )
        bone3d.append(diff3d)
        bone2d.append(diff2d)
        conf2d.append(minconf2d)

    bone3d = torch.stack(bone3d, dim=1)
    bone2d = torch.stack(bone2d, dim=2)
    bone2d = bone2d[np.arange(B), best_2d_idx, :]
    conf2d = torch.stack(conf2d, dim=2)
    conf2d = conf2d[np.arange(B), best_2d_idx, :]

    # mean over all bones
    mean_bone3d = torch.mean(bone3d, dim=1)
    mean_bone2d = torch.mean(bone2d * (conf2d > 0.0), dim=1)

    # approx z based on ratio
    init_z = focal * (mean_bone3d / mean_bone2d)
    init_trans = torch.zeros(B, T, 3)
    init_trans[:, :, 2] = init_z.unsqueeze(1).expand(B, T).detach()

    return init_trans
//  [python]
def estimate_initial_trans(body_pose, joints3d_op, joints2d_op, focal):
    """
    use focal length and bone lengths to approximate distance from camera
    (based on PROX https://github.com/mohamedhassanmus/prox/blob/master/prox/fitting.py)
    :param body_pose
    :param joints3d_op OP keypoints in 3d
    :param joints2d_op OP keypoints in 2d
    """
    B, T, N, _ = joints2d_op.shape
    joints2d_obs = joints2d_op[:, :, :, :2]  # (B, T, N, 3)
    joints2d_conf = joints2d_op[:, :, :, 2]

    # find least-occluded 2d frame
    num_2d_vis = torch.sum(joints2d_conf > 0.0, dim=2)
    best_2d_idx = torch.max(num_2d_vis, dim=1)[1]

    # calculate bone lengths and confidence in each bone length
    bone3d = []
    bone2d = []
    conf2d = []
    for pair in OP_EDGE_LIST:
        diff3d = torch.linalg.norm(
            joints3d_op[:, 0, pair[0], :] - joints3d_op[:, 0, pair[1], :],
            dim=1,
        )  # does not change over time
        diff2d = torch.linalg.norm(
            joints2d_obs[:, :, pair[0], :] - joints2d_obs[:, :, pair[1], :],
            dim=2,
        )
        minconf2d = torch.min(
            joints2d_conf[:, :, pair[0]], joints2d_conf[:, :, pair[1]]
        )
        bone3d.append(diff3d)
        bone2d.append(diff2d)
        conf2d.append(minconf2d)

    bone3d = torch.stack(bone3d, dim=1)
    bone2d = torch.stack(bone2d, dim=2)
    bone2d = bone2d[np.arange(B), best_2d_idx, :]
    conf2d = torch.stack(conf2d, dim=2)
    conf2d = conf2d[np.arange(B), best_2d_idx, :]

    # mean over all bones
    mean_bone3d = torch.mean(bone3d, dim=1)
    mean_bone2d = torch.mean(bone2d * (conf2d > 0.0), dim=1)

    # approx z based on ratio
    init_z = focal * (mean_bone3d / mean_bone2d)
    init_trans = torch.zeros(B, T, 3)
    init_trans[:, :, 2] = init_z.unsqueeze(1).expand(B, T).detach()

    return init_trans
//  [python]
def compute_world2prior(floor_plane, trans, root_orient, origin):
    """
    Computes rotation and translation from the camera frame to the canonical coordinate system
    used by the motion and initial state priors.
    - floor_plane : B x 3
    - trans : B x 3
    - root_orient : B x 3
    - origin: B x 3 desired origin (first joint of human model)
    - returns R (B, 3, 3), t (B, 3), root_height (B)
    """
    B = trans.size(0)
    if floor_plane.size(1) == 3:
        floor_plane_4d = parse_floor_plane(floor_plane)
    else:
        floor_plane_4d = floor_plane
    floor_normal = floor_plane_4d[:, :3]
    floor_trans, _ = compute_plane_intersection(trans, -floor_normal, floor_plane_4d)

    # compute prior frame axes within the camera frame
    # up is the floor_plane normal
    up_axis = floor_normal
    # right is body -x direction projected to floor plane
    root_orient_mat = batch_rodrigues(root_orient)
    body_right = -root_orient_mat[:, :, 0]
    floor_body_right, s = compute_plane_intersection(trans, body_right, floor_plane_4d)
    right_axis = floor_body_right - floor_trans
    # body right may not actually intersect - in this case must negate axis because we have the -x
    right_axis = torch.where(s.reshape((B, 1)) < 0, -right_axis, right_axis)
    right_axis = right_axis / torch.norm(right_axis, dim=1, keepdim=True)
    # forward is their cross product
    fwd_axis = torch.linalg.cross(up_axis, right_axis)
    fwd_axis = fwd_axis / torch.norm(fwd_axis, dim=1, keepdim=True)

    # prior frame is right, fwd, up
    prior_R = torch.stack([right_axis, fwd_axis, up_axis], dim=2)
    world2prior_R = prior_R.transpose(-1, -2)

    # translation takes translation to origin plus offset to the floor
    world2prior_t = -trans

    # compute the distance from root to ground plane
    _, s_root = compute_plane_intersection(origin, -floor_normal, floor_plane_4d)
    root_height = s_root.reshape(B, 1)

    if False:
        # apply transform to origin
        origin_prior = torch.matmul(world2prior_R, (origin - trans)[..., None])[..., 0]
        root_height = s_root.reshape(B, 1) - origin_prior[:, 2:3]

    return world2prior_R, world2prior_t, root_height
//  [python]
def compute_world2prior(floor_plane, trans, root_orient, origin):
    """
    Computes rotation and translation from the camera frame to the canonical coordinate system
    used by the motion and initial state priors.
    - floor_plane : B x 3
    - trans : B x 3
    - root_orient : B x 3
    - origin: B x 3 desired origin (first joint of human model)
    - returns R (B, 3, 3), t (B, 3), root_height (B)
    """
    B = trans.size(0)
    if floor_plane.size(1) == 3:
        floor_plane_4d = parse_floor_plane(floor_plane)
    else:
        floor_plane_4d = floor_plane
    floor_normal = floor_plane_4d[:, :3]
    floor_trans, _ = compute_plane_intersection(trans, -floor_normal, floor_plane_4d)

    # compute prior frame axes within the camera frame
    # up is the floor_plane normal
    up_axis = floor_normal
    # right is body -x direction projected to floor plane
    root_orient_mat = batch_rodrigues(root_orient)
    body_right = -root_orient_mat[:, :, 0]
    floor_body_right, s = compute_plane_intersection(trans, body_right, floor_plane_4d)
    right_axis = floor_body_right - floor_trans
    # body right may not actually intersect - in this case must negate axis because we have the -x
    right_axis = torch.where(s.reshape((B, 1)) < 0, -right_axis, right_axis)
    right_axis = right_axis / torch.linalg.norm(right_axis, dim=1, keepdim=True)
    # forward is their cross product
    fwd_axis = torch.linalg.cross(up_axis, right_axis)
    fwd_axis = fwd_axis / torch.linalg.norm(fwd_axis, dim=1, keepdim=True)

    # prior frame is right, fwd, up
    prior_R = torch.stack([right_axis, fwd_axis, up_axis], dim=2)
    world2prior_R = prior_R.transpose(-1, -2)

    # translation takes translation to origin plus offset to the floor
    world2prior_t = -trans

    # compute the distance from root to ground plane
    _, s_root = compute_plane_intersection(origin, -floor_normal, floor_plane_4d)
    root_height = s_root.reshape(B, 1)

    if False:
        # apply transform to origin
        origin_prior = torch.matmul(world2prior_R, (origin - trans)[..., None])[..., 0]
        root_height = s_root.reshape(B, 1) - origin_prior[:, 2:3]

    return world2prior_R, world2prior_t, root_height
//  [python]
def forward(ctx, input):
        x = input
        batchSize = x.data.shape[0]
        dim = x.data.shape[1]
        dtype = x.dtype
        normx = torch.norm(x, dim=[1, 2]).view(x.size(0), 1, 1)
        I = torch.eye(x.size(1), x.size(2), requires_grad=False, device=x.device).reshape(1, x.size(1),x.size(2)).repeat(x.size(0),1, 1).type(dtype)
        y = pade_approximant(x / normx, I)
        y = y * torch.sqrt(normx)
        ctx.save_for_backward(input, y, normx, I)
        return y
//  [python]
def forward(ctx, input):
        x = input
        batchSize = x.data.shape[0]
        dim = x.data.shape[1]
        dtype = x.dtype
        normx = torch.linalg.norm(x, dim=[1, 2]).view(x.size(0), 1, 1)
        I = torch.eye(x.size(1), x.size(2), requires_grad=False, device=x.device).reshape(1, x.size(1),x.size(2)).repeat(x.size(0),1, 1).type(dtype)
        y = pade_approximant(x / normx, I)
        y = y * torch.sqrt(normx)
        ctx.save_for_backward(input, y, normx, I)
        return y
//  [python]
def get_rays(H, W, c2i, w2c):
    i, j = torch.meshgrid(torch.linspace(0, H - 1, H).to(w2c.device),
                          torch.linspace(0, W - 1, W).to(w2c.device))
    i = i.t()
    j = j.t()

    cx, cy = (H - 1) / 2, (W - 1) / 2

    i2c = torch.linalg.pinv(c2i)

    c2w = torch.linalg.inv(w2c)
    ori = c2w @ torch.Tensor([0, 0, 0, 1]).to(w2c.device)
    ori = ori[:3]

    coords = torch.stack([i - cx,
                          j - cy,
                          torch.ones_like(i)], dim=1)
    rays_d = i2c @ coords
    rays_d = c2w @ rays_d
    rays_d = torch.moveaxis(rays_d, 1, -1)
    rays_d = rays_d[:, :, :3]
    rays_d /= torch.norm(rays_d, dim=-1, keepdim=True)

    rays_o = ori.expand(rays_d.shape)
    return rays_o.numpy(), rays_d.numpy()
//  [python]
def get_rays(H, W, c2i, w2c):
    i, j = torch.meshgrid(torch.linspace(0, H - 1, H).to(w2c.device),
                          torch.linspace(0, W - 1, W).to(w2c.device))
    i = i.t()
    j = j.t()

    cx, cy = (H - 1) / 2, (W - 1) / 2

    i2c = torch.linalg.pinv(c2i)

    c2w = torch.linalg.inv(w2c)
    ori = c2w @ torch.Tensor([0, 0, 0, 1]).to(w2c.device)
    ori = ori[:3]

    coords = torch.stack([i - cx,
                          j - cy,
                          torch.ones_like(i)], dim=1)
    rays_d = i2c @ coords
    rays_d = c2w @ rays_d
    rays_d = torch.moveaxis(rays_d, 1, -1)
    rays_d = rays_d[:, :, :3]
    rays_d /= torch.linalg.norm(rays_d, dim=-1, keepdim=True)

    rays_o = ori.expand(rays_d.shape)
    return rays_o.numpy(), rays_d.numpy()
//  [python]
def logit_excluded_and_restricted_loss(self, logits, sim, cube):
        """ 
        Restrict at the logit level by excluding and restricting to individual characters - this metric is not used in the paper
        """
        centered_logits = logits - logits.mean(dim=-1, keepdim=True)
        centered_logits = centered_logits.reshape(-1)
        trace = cube.reshape(-1)
        unit_trace = trace / torch.norm(trace)
        norm_logits = torch.norm(centered_logits)
        excluded_logits = centered_logits - sim * norm_logits * unit_trace
        restricted_logits = sim * norm_logits * unit_trace
        excluded_logits = excluded_logits.reshape(-1, self.group.order)
        excluded_logits = excluded_logits[self.train_indices]
        excluded_loss = loss_fn(excluded_logits, self.train_labels).item()
        restricted_loss = loss_fn(restricted_logits.reshape(-1, self.group.order), self.all_labels).item()
        return excluded_loss, restricted_loss
//  [python]
def logit_excluded_and_restricted_loss(self, logits, sim, cube):
        """ 
        Restrict at the logit level by excluding and restricting to individual characters - this metric is not used in the paper
        """
        centered_logits = logits - logits.mean(dim=-1, keepdim=True)
        centered_logits = centered_logits.reshape(-1)
        trace = cube.reshape(-1)
        unit_trace = trace / torch.linalg.norm(trace)
        norm_logits = torch.linalg.norm(centered_logits)
        excluded_logits = centered_logits - sim * norm_logits * unit_trace
        restricted_logits = sim * norm_logits * unit_trace
        excluded_logits = excluded_logits.reshape(-1, self.group.order)
        excluded_logits = excluded_logits[self.train_indices]
        excluded_loss = loss_fn(excluded_logits, self.train_labels).item()
        restricted_loss = loss_fn(restricted_logits.reshape(-1, self.group.order), self.all_labels).item()
        return excluded_loss, restricted_loss
//  [python]
def total_logit_excluded_and_restricted_loss(self, logits, key_reps, sims):
        """ 
        Restrict at the logit level by excluding and restricting to key characters - this metric is not used in the paper
        """
        centered_logits = logits - logits.mean(dim=-1, keepdim=True)
        centered_logits = centered_logits.reshape(-1)
        norm_logits = torch.norm(centered_logits)
        
        key_rep_logits = torch.zeros_like(centered_logits)
        for rep_name in key_reps:
            trace = self.group.irreps[rep_name].logit_trace_tensor_cube.reshape(-1)
            unit_trace = trace / torch.norm(trace)

            key_rep_logits += sims[rep_name] * norm_logits * unit_trace

        excluded_logits = centered_logits - key_rep_logits
        excluded_logits = excluded_logits.reshape(-1, self.group.order)
        excluded_logits = excluded_logits[self.train_indices]
        restricted_logits = key_rep_logits

        excluded_loss = loss_fn(excluded_logits, self.train_labels).item()
        restricted_loss = loss_fn(restricted_logits.reshape(-1, self.group.order), self.all_labels).item()

        return excluded_loss, restricted_loss
//  [python]
def total_logit_excluded_and_restricted_loss(self, logits, key_reps, sims):
        """ 
        Restrict at the logit level by excluding and restricting to key characters - this metric is not used in the paper
        """
        centered_logits = logits - logits.mean(dim=-1, keepdim=True)
        centered_logits = centered_logits.reshape(-1)
        norm_logits = torch.linalg.norm(centered_logits)
        
        key_rep_logits = torch.zeros_like(centered_logits)
        for rep_name in key_reps:
            trace = self.group.irreps[rep_name].logit_trace_tensor_cube.reshape(-1)
            unit_trace = trace / torch.linalg.norm(trace)

            key_rep_logits += sims[rep_name] * norm_logits * unit_trace

        excluded_logits = centered_logits - key_rep_logits
        excluded_logits = excluded_logits.reshape(-1, self.group.order)
        excluded_logits = excluded_logits[self.train_indices]
        restricted_logits = key_rep_logits

        excluded_loss = loss_fn(excluded_logits, self.train_labels).item()
        restricted_loss = loss_fn(restricted_logits.reshape(-1, self.group.order), self.all_labels).item()

        return excluded_loss, restricted_loss
//  [python]
def joints2bones(joints, parents):
    ''' Decompose joints location to bone length and direction.

        Parameters
        ----------
        joints: torch.tensor Bx24x3
    '''
    assert joints.shape[1] == parents.shape[0]
    bone_dirs = torch.zeros_like(joints)
    bone_lens = torch.zeros_like(joints[:, :, :1])

    for c_id in range(parents.shape[0]):
        p_id = parents[c_id]
        if p_id == -1:
            # Parent node
            bone_dirs[:, c_id] = joints[:, c_id]
        else:
            # Child node
            # (B, 3)
            diff = joints[:, c_id] - joints[:, p_id]
            length = torch.norm(diff, dim=1, keepdim=True) + 1e-8
            direct = diff / length

            bone_dirs[:, c_id] = direct
            bone_lens[:, c_id] = length

    return bone_dirs, bone_lens
//  [python]
def joints2bones(joints, parents):
    ''' Decompose joints location to bone length and direction.

        Parameters
        ----------
        joints: torch.tensor Bx24x3
    '''
    assert joints.shape[1] == parents.shape[0]
    bone_dirs = torch.zeros_like(joints)
    bone_lens = torch.zeros_like(joints[:, :, :1])

    for c_id in range(parents.shape[0]):
        p_id = parents[c_id]
        if p_id == -1:
            # Parent node
            bone_dirs[:, c_id] = joints[:, c_id]
        else:
            # Child node
            # (B, 3)
            diff = joints[:, c_id] - joints[:, p_id]
            length = torch.linalg.norm(diff, dim=1, keepdim=True) + 1e-8
            direct = diff / length

            bone_dirs[:, c_id] = direct
            bone_lens[:, c_id] = length

    return bone_dirs, bone_lens
//  [python]
def vanishing_point(X):
    l1 = X[..., 0, 6:9]
    l2 = X[..., 1, 6:9]

    vp = torch.cross(l1, l2, dim=-1)

    vp_norm = torch.norm(vp, keepdim=True, dim=-1) + 1e-8
    vp = vp / vp_norm

    return vp
//  [python]
def vanishing_point(X):
    l1 = X[..., 0, 6:9]
    l2 = X[..., 1, 6:9]

    vp = torch.cross(l1, l2, dim=-1)

    vp_norm = torch.linalg.norm(vp, keepdim=True, dim=-1) + 1e-8
    vp = vp / vp_norm

    return vp
//  [python]
def slerp(val, low, high):
    low_norm = low/torch.norm(low, dim=1, keepdim=True)
    high_norm = high/torch.norm(high, dim=1, keepdim=True)
    dot = (low_norm*high_norm).sum(1)

    if dot.mean() > 0.9995:
        return low * val + high * (1 - val)

    omega = torch.acos(dot)
    so = torch.sin(omega)
    res = (torch.sin((1.0-val)*omega)/so).unsqueeze(1)*low + (torch.sin(val*omega)/so).unsqueeze(1) * high
    return res
//  [python]
def slerp(val, low, high):
    low_norm = low/torch.linalg.norm(low, dim=1, keepdim=True)
    high_norm = high/torch.linalg.norm(high, dim=1, keepdim=True)
    dot = (low_norm*high_norm).sum(1)

    if dot.mean() > 0.9995:
        return low * val + high * (1 - val)

    omega = torch.acos(dot)
    so = torch.sin(omega)
    res = (torch.sin((1.0-val)*omega)/so).unsqueeze(1)*low + (torch.sin(val*omega)/so).unsqueeze(1) * high
    return res
//  [python]
def forward(self, states):
        N, state_dim = states.shape
        
        quat = states[:, 6:10].detach().clone()
        quat = torch.div(quat, torch.norm(quat, dim=1)[:, None])
        
        vel_enu = states[:, 3:6]*self.scale
        
        body_velocity = tf.quaternion_apply(quat, vel_enu)
#         print(vel_enu[0].detach().numpy(), body_velocity[0].detach().numpy(), orientation[0].detach().numpy())
        
        R = torch.diag(torch.stack([torch.tensor(1e-2), self.std, torch.tensor(1e-2)])).expand((N, self.observation_dim, self.observation_dim))
        
        return body_velocity.float(), R.float()
//  [python]
def forward(self, states):
        N, state_dim = states.shape
        
        quat = states[:, 6:10].detach().clone()
        quat = torch.div(quat, torch.linalg.norm(quat, dim=1)[:, None])
        
        vel_enu = states[:, 3:6]*self.scale
        
        body_velocity = tf.quaternion_apply(quat, vel_enu)
#         print(vel_enu[0].detach().numpy(), body_velocity[0].detach().numpy(), orientation[0].detach().numpy())
        
        R = torch.diag(torch.stack([torch.tensor(1e-2), self.std, torch.tensor(1e-2)])).expand((N, self.observation_dim, self.observation_dim))
        
        return body_velocity.float(), R.float()