jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def getDistances(control_w):
    p1 = control_w[:3,0]
    p2 = control_w[:3,1]
    p3 = control_w[:3,2]
    p4 = control_w[:3,3]

    d12 = torch.linalg.norm(p1-p2,ord=2)
    d13 = torch.linalg.norm(p1-p3,ord=2)
    d14 = torch.linalg.norm(p1-p4,ord=2)
    d23 = torch.linalg.norm(p2-p3,ord=2)
    d24 = torch.linalg.norm(p2-p4,ord=2)
    d34 = torch.linalg.norm(p3-p4,ord=2)

    return d12,d13,d14,d23,d24,d34
//  [python]
def getError2(pimg,pcam,A,show=False):
    M = pimg.shape[0]
    N = pimg.shape[1]

    proj = torch.bmm(torch.stack(M*[A]),pcam)
    proj = proj_img = proj / proj[:,-1,:].unsqueeze(1)
    pimg_pred = proj_img[:,:2,:]
    diff = pimg - pimg_pred
    error = torch.mean(torch.norm(diff,p=2,dim=1),dim=1)

    return error
//  [python]
def getError2(pimg,pcam,A,show=False):
    M = pimg.shape[0]
    N = pimg.shape[1]

    proj = torch.bmm(torch.stack(M*[A]),pcam)
    proj = proj_img = proj / proj[:,-1,:].unsqueeze(1)
    pimg_pred = proj_img[:,:2,:]
    diff = pimg - pimg_pred
    error = torch.mean(torch.linalg.norm(diff,p=2,dim=1),dim=1)

    return error
//  [python]
def getShapeError(s1,s2):
    diff = s1 - s2
    error = torch.mean(torch.norm(diff,dim=1))
    return error
//  [python]
def getShapeError(s1,s2):
    diff = s1 - s2
    error = torch.mean(torch.linalg.norm(diff,dim=1))
    return error
//  [python]
def scaleControlPoints_single(c_c,c_w,alphas,x_w):
    x_c = torch.mm(c_c,alphas)

    centered_xw = x_w - torch.mean(x_w,dim=0).unsqueeze(0)
    d_w = torch.norm(centered_xw,p=2,dim=1)
    centered_xc = x_c - torch.mean(x_c,dim=1).unsqueeze(1)
    d_c = torch.norm(centered_xc,p=2,dim=0)

    # least square solution to scale
    s = 1.0 /((1.0/torch.sum(d_c*d_c) * torch.sum(d_c*d_w)))

    # apply scale onto c_c and recompute the camera coordinates
    sc_c = c_c / s
    sx_c = torch.mm(sc_c,alphas).T

    # fix the sign so negative depth is not possible
    negdepth_mask = sx_c[:,-1] < 0
    sx_c = sx_c * (negdepth_mask.float().unsqueeze(1) * -2 +1)

    return sc_c, sx_c, s
//  [python]
def scaleControlPoints_single(c_c,c_w,alphas,x_w):
    x_c = torch.mm(c_c,alphas)

    centered_xw = x_w - torch.mean(x_w,dim=0).unsqueeze(0)
    d_w = torch.linalg.norm(centered_xw,p=2,dim=1)
    centered_xc = x_c - torch.mean(x_c,dim=1).unsqueeze(1)
    d_c = torch.linalg.norm(centered_xc,p=2,dim=0)

    # least square solution to scale
    s = 1.0 /((1.0/torch.sum(d_c*d_c) * torch.sum(d_c*d_w)))

    # apply scale onto c_c and recompute the camera coordinates
    sc_c = c_c / s
    sx_c = torch.mm(sc_c,alphas).T

    # fix the sign so negative depth is not possible
    negdepth_mask = sx_c[:,-1] < 0
    sx_c = sx_c * (negdepth_mask.float().unsqueeze(1) * -2 +1)

    return sc_c, sx_c, s
//  [python]
def getReprojError3(xcam,pw,R,T):

    M = xcam.shape[0]
    pc = torch.bmm(R,torch.stack(M*[pw]).permute(0,2,1))
    pct = pc + T.unsqueeze(2)
    diff = pct - xcam
    error  = torch.mean(torch.norm(diff,p=2,dim=1),dim=1)

    return error
//  [python]
def getReprojError3(xcam,pw,R,T):

    M = xcam.shape[0]
    pc = torch.bmm(R,torch.stack(M*[pw]).permute(0,2,1))
    pct = pc + T.unsqueeze(2)
    diff = pct - xcam
    error  = torch.mean(torch.linalg.norm(diff,ord=2,dim=1),dim=1)

    return error
//  [python]
def scaleControlPoints(c_c,c_w,alphas,x_w):
    views = c_c.shape[0]
    rep_alpha = torch.stack(views*[alphas])
    x_c = torch.bmm(c_c,rep_alpha)

    # center x_w and x_c
    centered_xw = x_w - torch.mean(x_w,dim=0).unsqueeze(0)
    d_w = torch.norm(centered_xw,p=2,dim=1)
    centered_xc = x_c - torch.mean(x_c,dim=2).unsqueeze(2)
    d_c = torch.norm(centered_xc,p=2,dim=1)

    # least square solution to scale
    s = 1.0 /((1.0/(torch.sum(d_c*d_c,dim=1)) * torch.sum(d_c*d_w.unsqueeze(0),dim=1)))

    # apply scale onto c_c and recompute the camera coordinates
    sc_c = c_c / s.unsqueeze(1).unsqueeze(1)
    sx_c = torch.bmm(sc_c,rep_alpha).permute(0,2,1)

    # fix the sign so negative depth is not possible
    #print(sx_c.shape)
    negdepth_mask = sx_c[:,:,-1] < 0
    sx_c = sx_c * (negdepth_mask.float().unsqueeze(2) * -2 +1)

    return sc_c, sx_c, s
//  [python]
def scaleControlPoints(c_c,c_w,alphas,x_w):
    views = c_c.shape[0]
    rep_alpha = torch.stack(views*[alphas])
    x_c = torch.bmm(c_c,rep_alpha)

    # center x_w and x_c
    centered_xw = x_w - torch.mean(x_w,dim=0).unsqueeze(0)
    d_w = torch.linalg.norm(centered_xw,p=2,dim=1)
    centered_xc = x_c - torch.mean(x_c,dim=2).unsqueeze(2)
    d_c = torch.linalg.norm(centered_xc,p=2,dim=1)

    # least square solution to scale
    s = 1.0 /((1.0/(torch.sum(d_c*d_c,dim=1)) * torch.sum(d_c*d_w.unsqueeze(0),dim=1)))

    # apply scale onto c_c and recompute the camera coordinates
    sc_c = c_c / s.unsqueeze(1).unsqueeze(1)
    sx_c = torch.bmm(sc_c,rep_alpha).permute(0,2,1)

    # fix the sign so negative depth is not possible
    #print(sx_c.shape)
    negdepth_mask = sx_c[:,:,-1] < 0
    sx_c = sx_c * (negdepth_mask.float().unsqueeze(2) * -2 +1)

    return sc_c, sx_c, s
//  [python]
def scaleControlPoints_(c_c,c_w,alphas,x_w):
    views = c_c.shape[0]
    x_c = torch.bmm(c_c,alphas)

    # center x_w and x_c
    centered_xw = x_w - torch.mean(x_w,dim=1).unsqueeze(1)

    # get depth to world points
    d_w = torch.norm(centered_xw,dim=2)

    # get depth to control points
    centered_xc = x_c - torch.mean(x_c,dim=2).unsqueeze(2)
    d_c = torch.norm(centered_xc,dim=1)

    # least square solution to scale
    #s = 1.0 /((1.0/(torch.sum(d_c*d_c,dim=1)) * torch.sum(d_c*d_w,dim=1)))
    #s = 1.0 /( torch.sum(d_c*d_w,dim=1)/ torch.sum(d_c*d_c,dim=1))
    s = torch.sum(d_c*d_w,dim=1) / torch.sum(d_c*d_c,dim=1)

    # apply scale onto c_c and recompute the camera coordinates
    sc_c = c_c * s.unsqueeze(1).unsqueeze(1)
    sx_c = torch.bmm(sc_c,alphas).permute(0,2,1)

    # fix the sign so negative depth is not possible
    negdepth_mask = sx_c[:,:,-1] < 0
    # sx_c = sx_c * (negdepth_mask.float().unsqueeze(2) * -2 +1)

    # scale control points
    #mask = sc_c[:,:,-1] < 0
    #print(torch.any(negdepth_mask))
    #quit()

    return s
//  [python]
def scaleControlPoints_(c_c,c_w,alphas,x_w):
    views = c_c.shape[0]
    x_c = torch.bmm(c_c,alphas)

    # center x_w and x_c
    centered_xw = x_w - torch.mean(x_w,dim=1).unsqueeze(1)

    # get depth to world points
    d_w = torch.linalg.norm(centered_xw,dim=2)

    # get depth to control points
    centered_xc = x_c - torch.mean(x_c,dim=2).unsqueeze(2)
    d_c = torch.linalg.norm(centered_xc,dim=1)

    # least square solution to scale
    #s = 1.0 /((1.0/(torch.sum(d_c*d_c,dim=1)) * torch.sum(d_c*d_w,dim=1)))
    #s = 1.0 /( torch.sum(d_c*d_w,dim=1)/ torch.sum(d_c*d_c,dim=1))
    s = torch.sum(d_c*d_w,dim=1) / torch.sum(d_c*d_c,dim=1)

    # apply scale onto c_c and recompute the camera coordinates
    sc_c = c_c * s.unsqueeze(1).unsqueeze(1)
    sx_c = torch.bmm(sc_c,alphas).permute(0,2,1)

    # fix the sign so negative depth is not possible
    negdepth_mask = sx_c[:,:,-1] < 0
    # sx_c = sx_c * (negdepth_mask.float().unsqueeze(2) * -2 +1)

    # scale control points
    #mask = sc_c[:,:,-1] < 0
    #print(torch.any(negdepth_mask))
    #quit()

    return s
//  [python]
def getBatchError(pts,shape,K):
    e2d = []
    e3d = []
    e2d_all = []
    e3d_all = []
    d_all = []
    b = pts.shape[0]
    for i in range(b):
        pW = shape[i]
        pI = pts[i]
        M = pI.shape[0]
        km, c_w, scaled_betas, alphas = EPnP(pI,pW,K)

        Xc, R, T, mask = optimizeGN(km,c_w,scaled_betas,alphas,pW,pI,K)

        error2d = getReprojError2(pI,pW,R,T,K,loss='l2')
        error3d = getRelReprojError3(pC,pW,R,T)
        d = torch.norm(pC,p=2,dim=1).mean(1)

        e2d.append(error2d.mean())
        e3d.append(error3d.mean())
        e2d_all.append(error2d)
        e3d_all.append(error3d)
        d_all.append(d)
    errorShape = torch.mean(torch.norm(xw_gt - shape,dim=2),dim=1)
    return torch.stack(e2d),torch.stack(e3d),errorShape,torch.stack(e2d_all),torch.stack(e3d_all),torch.stack(d_all)
//  [python]
def getBatchError(pts,shape,K):
    e2d = []
    e3d = []
    e2d_all = []
    e3d_all = []
    d_all = []
    b = pts.shape[0]
    for i in range(b):
        pW = shape[i]
        pI = pts[i]
        M = pI.shape[0]
        km, c_w, scaled_betas, alphas = EPnP(pI,pW,K)

        Xc, R, T, mask = optimizeGN(km,c_w,scaled_betas,alphas,pW,pI,K)

        error2d = getReprojError2(pI,pW,R,T,K,loss='l2')
        error3d = getRelReprojError3(pC,pW,R,T)
        d = torch.linalg.norm(pC,p=2,dim=1).mean(1)

        e2d.append(error2d.mean())
        e3d.append(error3d.mean())
        e2d_all.append(error2d)
        e3d_all.append(error3d)
        d_all.append(d)
    errorShape = torch.mean(torch.linalg.norm(xw_gt - shape,dim=2),dim=1)
    return torch.stack(e2d),torch.stack(e3d),errorShape,torch.stack(e2d_all),torch.stack(e3d_all),torch.stack(d_all)
//  [python]
def getTimeConsistency2(pW,R,T):
    M = R.shape[0]
    pc = torch.bmm(R,torch.stack(M*[pW]).permute(0,2,1))
    pct = pc + T.unsqueeze(2)

    error = torch.norm(pct[:M-1,:,:] - pct[1:,:,:],dim=1)
    return error
//  [python]
def getTimeConsistency2(pW,R,T):
    M = R.shape[0]
    pc = torch.bmm(R,torch.stack(M*[pW]).permute(0,2,1))
    pct = pc + T.unsqueeze(2)

    error = torch.linalg.norm(pct[:M-1,:,:] - pct[1:,:,:],dim=1)
    return error
//  [python]
def getTimeConsistency(pW,R,T):
    M = R.shape[0]
    pc = torch.bmm(R,torch.stack(M*[pW]).permute(0,2,1))
    pct = pc + T.unsqueeze(2)

    error = torch.mean(torch.norm(pct[:M-1,:,:] - pct[1:,:,:],dim=1))
    return error
//  [python]
def getTimeConsistency(pW,R,T):
    M = R.shape[0]
    pc = torch.bmm(R,torch.stack(M*[pW]).permute(0,2,1))
    pct = pc + T.unsqueeze(2)

    error = torch.mean(torch.linalg.norm(pct[:M-1,:,:] - pct[1:,:,:],dim=1))
    return error
//  [python]
def getRelReprojError3(xcam,pw,R,T):

    M = xcam.shape[0]
    pc = torch.bmm(R,torch.stack(M*[pw]).permute(0,2,1))
    pct = pc + T.unsqueeze(2)

    dgt = torch.norm(torch.mean(pct,dim=2),dim=1)
    d = torch.norm(torch.mean(xcam,dim=2),dim=1)
    diff = torch.abs(dgt - d)

    return diff / d
//  [python]
def getRelReprojError3(xcam,pw,R,T):

    M = xcam.shape[0]
    pc = torch.bmm(R,torch.stack(M*[pw]).permute(0,2,1))
    pct = pc + T.unsqueeze(2)

    dgt = torch.linalg.norm(torch.mean(pct,dim=2),dim=1)
    d = torch.linalg.norm(torch.mean(xcam,dim=2),dim=1)
    diff = torch.abs(dgt - d)

    return diff / d
//  [python]
def apply(self, b_mat, a_mat, c_cur):
        if c_cur is None:
            return torch.ones((a_mat.size(0), ), device=a_mat.device)
        fit_res = torch.norm(b_mat - torch.matmul(a_mat, c_cur), dim=-1)
        sigma = self.robust_k
        w_func = torch.where(fit_res < sigma, torch.ones_like(fit_res), sigma / fit_res)
        return torch.sqrt(w_func)
//  [python]
def apply(self, b_mat, a_mat, c_cur):
        if c_cur is None:
            return torch.ones((a_mat.size(0), ), device=a_mat.device)
        fit_res = torch.linalg.norm(b_mat - torch.matmul(a_mat, c_cur), dim=-1)
        sigma = self.robust_k
        w_func = torch.where(fit_res < sigma, torch.ones_like(fit_res), sigma / fit_res)
        return torch.sqrt(w_func)
//  [python]
def quat2euler(quat):

    dist = torch.norm(quat,p=2,dim=0)
    q = quat / dist.unsqueeze(0)
    rx = torch.atan2(2*(q[0]*q[1] + q[2]*q[3]),1-2*(q[1]**2+q[2]**2));
    ry = torch.asin(2*(q[0]*q[2] - q[3]*q[1]));
    rz = torch.atan2(2*(q[0]*q[3] + q[1]*q[2]),1-2*(q[2]**2+q[3]**2));

    return rx, ry,rz
//  [python]
def quat2euler(quat):

    dist = torch.linalg.norm(quat,ord=2,dim=0)
    q = quat / dist.unsqueeze(0)
    rx = torch.atan2(2*(q[0]*q[1] + q[2]*q[3]),1-2*(q[1]**2+q[2]**2));
    ry = torch.asin(2*(q[0]*q[2] - q[3]*q[1]));
    rz = torch.atan2(2*(q[0]*q[3] + q[1]*q[2]),1-2*(q[2]**2+q[3]**2));

    return rx, ry,rz
//  [python]
def get3DConsistency(pI,pW,kinv,R,T):
    kinv[0,0] = 1/100
    kinv[1,1] = 1/100
    M = pI.shape[0]
    N = pI.shape[2]

    pC = torch.bmm(R,torch.stack(M*[pW.T])) + T.unsqueeze(2)
    z = pC[:,2,:]
    ones = torch.ones(M,1,N).to(pI.device)

    pI_h = torch.cat((pI,ones),dim=1)
    pI_proj = torch.bmm(torch.stack(M*[kinv]),pI_h)
    pC_proj = pI_proj * z.unsqueeze(1)

    le_gt = torch.mean(pW[36:42,:],dim=0)
    re_gt = torch.mean(pW[42:48,:],dim=0)
    d_gt = torch.norm(le_gt - re_gt)

    le = torch.mean(pC_proj[:,:,36:42],dim=2)
    re = torch.mean(pC_proj[:,:,42:48],dim=2)
    d = torch.norm(le - re,dim=1)

    return torch.mean(torch.abs(d - d_gt))
//  [python]
def get3DConsistency(pI,pW,kinv,R,T):
    kinv[0,0] = 1/100
    kinv[1,1] = 1/100
    M = pI.shape[0]
    N = pI.shape[2]

    pC = torch.bmm(R,torch.stack(M*[pW.T])) + T.unsqueeze(2)
    z = pC[:,2,:]
    ones = torch.ones(M,1,N).to(pI.device)

    pI_h = torch.cat((pI,ones),dim=1)
    pI_proj = torch.bmm(torch.stack(M*[kinv]),pI_h)
    pC_proj = pI_proj * z.unsqueeze(1)

    le_gt = torch.mean(pW[36:42,:],dim=0)
    re_gt = torch.mean(pW[42:48,:],dim=0)
    d_gt = torch.linalg.norm(le_gt - re_gt)

    le = torch.mean(pC_proj[:,:,36:42],dim=2)
    re = torch.mean(pC_proj[:,:,42:48],dim=2)
    d = torch.linalg.norm(le - re,dim=1)

    return torch.mean(torch.abs(d - d_gt))
//  [python]
def kernel_normalize(self, kernel, k=3):
        b, ch, h, w, kk = kernel.size()
        
        # calc kernel norm
        kernel = kernel.view(b, ch, h*w, kk).transpose(2, 1)
        kernel_norm = torch.norm(kernel.contiguous().view(b, h*w, ch*kk), p=2, dim=2, keepdim=True)
        
        # kernel reshape
        kernel = kernel.view(b, h*w, ch, k, k)
        kernel_norm = kernel_norm.view(b, h*w, 1, 1, 1)
        
        return kernel, kernel_norm
//  [python]
def kernel_normalize(self, kernel, k=3):
        b, ch, h, w, kk = kernel.size()
        
        # calc kernel norm
        kernel = kernel.view(b, ch, h*w, kk).transpose(2, 1)
        kernel_norm = torch.linalg.norm(kernel.contiguous().view(b, h*w, ch*kk), ord=2, dim=2, keepdim=True)
        
        # kernel reshape
        kernel = kernel.view(b, h*w, ch, k, k)
        kernel_norm = kernel_norm.view(b, h*w, 1, 1, 1)
        
        return kernel, kernel_norm
//  [python]
def getReprojError2_(pimg,pct,A,show=False,loss='l2'):

    M = pimg.shape[0]
    N = pimg.shape[2]
    proj = torch.bmm(torch.stack(M*[A]),pct)
    proj_img = proj / proj[:,-1,:].unsqueeze(1)

    pimg_pred = proj_img[:,:2,:]
    diff = pimg - pimg_pred
    if loss == 'l2':
        error  = torch.norm(diff,p=2,dim=1).mean(1)
    elif loss == 'l1':
        error = torch.abs(diff).mean(1).mean(1)

    if show:
        #import pptk
        #x = pct[-1].T.detach().cpu().numpy()
        #v = pptk.viewer(x)
        #v.set(point_size=1.1)
        #for i in range(M):
        #    pt1 = pimg[i].T.cpu().numpy()
        #    scatter(pt1)
        pta = pimg[0].T.cpu().numpy()
        ptb = pimg[-1].T.cpu().numpy()
        ptc = pimg_pred[0].detach().T.cpu().numpy()
        ptd = pimg_pred[-1].detach().T.cpu().numpy()
        plt.scatter(pta[:,0],pta[:,1],s=15,facecolors='none',edgecolors='green')
        #plt.scatter(ptb[:,0],pta[:,1],s=15,facecolors='none',edgecolors='green')
        plt.scatter(ptb[:,0],ptb[:,1],s=10,marker='.',edgecolors='red')
        #plt.scatter(ptb[:,0],pta[:,1],s=10,marker='.',edgecolors='red')

        #plt.xlim((-320,320))
        #plt.ylim((-240,240))
        plt.show()

        quit()

    return error
//  [python]
def getReprojError2_(pimg,pct,A,show=False,loss='l2'):

    M = pimg.shape[0]
    N = pimg.shape[2]
    proj = torch.bmm(torch.stack(M*[A]),pct)
    proj_img = proj / proj[:,-1,:].unsqueeze(1)

    pimg_pred = proj_img[:,:2,:]
    diff = pimg - pimg_pred
    if loss == 'l2':
        error  = torch.linalg.norm(diff,ord=2,dim=1).mean(1)
    elif loss == 'l1':
        error = torch.abs(diff).mean(1).mean(1)

    if show:
        #import pptk
        #x = pct[-1].T.detach().cpu().numpy()
        #v = pptk.viewer(x)
        #v.set(point_size=1.1)
        #for i in range(M):
        #    pt1 = pimg[i].T.cpu().numpy()
        #    scatter(pt1)
        pta = pimg[0].T.cpu().numpy()
        ptb = pimg[-1].T.cpu().numpy()
        ptc = pimg_pred[0].detach().T.cpu().numpy()
        ptd = pimg_pred[-1].detach().T.cpu().numpy()
        plt.scatter(pta[:,0],pta[:,1],s=15,facecolors='none',edgecolors='green')
        #plt.scatter(ptb[:,0],pta[:,1],s=15,facecolors='none',edgecolors='green')
        plt.scatter(ptb[:,0],ptb[:,1],s=10,marker='.',edgecolors='red')
        #plt.scatter(ptb[:,0],pta[:,1],s=10,marker='.',edgecolors='red')

        #plt.xlim((-320,320))
        #plt.ylim((-240,240))
        plt.show()

        quit()

    return error
//  [python]
def getReprojError2(pimg,pw,R,T,A,show=False,loss='l2'):

    M = pimg.shape[0]
    N = pimg.shape[2]
    pc = torch.bmm(R,torch.stack(M*[pw]).permute(0,2,1))
    pct = pc + T.unsqueeze(2)
    proj = torch.bmm(torch.stack(M*[A]),pct)
    proj_img = proj / proj[:,-1,:].unsqueeze(1)
    pimg_pred = proj_img[:,:2,:]

    #diff1 = pimg - pimg_pred
    #diff2 = pimg*-1 - pimg_pred
    #if loss == 'l2':
    #    error1 = torch.norm(diff1,p=2,dim=1)
    #    error2 = torch.norm(diff2,p=2,dim=1)
    #elif loss == 'l1':
    #    error1 = torch.abs(diff1)
    #    error2 = torch.abs(diff2)
    #error = error1 if torch.mean(error1) < torch.mean(error2) else error2

    error = torch.norm(pimg - pimg_pred,dim=1).mean()

    if show:
        for i in range(pimg.shape[0]):
            pta = pimg[i].T.cpu().numpy()
            ptc = pimg_pred[i].detach().T.cpu().numpy()

            plt.scatter(pta[:,0],pta[:,1],s=15,facecolors='none',edgecolors='green')
            plt.scatter(ptc[:,0],ptc[:,1],s=10,marker='.',edgecolors='red')
            plt.show()
            break

    return error
//  [python]
def getReprojError2(pimg,pw,R,T,A,show=False,loss='l2'):

    M = pimg.shape[0]
    N = pimg.shape[2]
    pc = torch.bmm(R,torch.stack(M*[pw]).permute(0,2,1))
    pct = pc + T.unsqueeze(2)
    proj = torch.bmm(torch.stack(M*[A]),pct)
    proj_img = proj / proj[:,-1,:].unsqueeze(1)
    pimg_pred = proj_img[:,:2,:]

    #diff1 = pimg - pimg_pred
    #diff2 = pimg*-1 - pimg_pred
    #if loss == 'l2':
    #    error1 = torch.linalg.norm(diff1,ord=2,dim=1)
    #    error2 = torch.linalg.norm(diff2,ord=2,dim=1)
    #elif loss == 'l1':
    #    error1 = torch.abs(diff1)
    #    error2 = torch.abs(diff2)
    #error = error1 if torch.mean(error1) < torch.mean(error2) else error2

    error = torch.linalg.norm(pimg - pimg_pred,ord=2,dim=1).mean()

    if show:
        for i in range(pimg.shape[0]):
            pta = pimg[i].T.cpu().numpy()
            ptc = pimg_pred[i].detach().T.cpu().numpy()

            plt.scatter(pta[:,0],pta[:,1],s=15,facecolors='none',edgecolors='green')
            plt.scatter(ptc[:,0],ptc[:,1],s=10,marker='.',edgecolors='red')
            plt.show()
            break

    return error
//  [python]
def forward(self, x):
        gx = torch.norm(x, p=2, dim=(1, 2), keepdim=True)
        nx = gx / (gx.mean(dim=-1, keepdim=True) + 1e-6)
        return self.gamma * (x * nx) + self.beta + x
//  [python]
def forward(self, x):
        gx = torch.linalg.norm(x, ord=2, dim=(1, 2), keepdim=True)
        nx = gx / (gx.mean(dim=-1, keepdim=True) + 1e-6)
        return self.gamma * (x * nx) + self.beta + x
//  [python]
def run(self, us, mr, radius=9, alpha=1e-3, beta=1e-2):
        us = us.squeeze(1)
        mr = mr.squeeze(1)

        bs = mr.size(0)
        pad = (mr.size(1) - (2*radius+1)) // 2
        count = (2*radius+1)**3

        self.f = self.f.to(mr.device)

        grad = torch.norm(F.conv3d(mr.unsqueeze(1), self.f, padding=1), dim=1)

        A = torch.ones(bs, 3, count, device=mr.device)
        A[:, 0] = mr[:, pad:-pad, pad:-pad, pad:-pad].reshape(bs, -1)
        A[:, 1] = grad[:, pad:-pad, pad:-pad, pad:-pad].reshape(bs, -1)
        b = us[:, pad:-pad, pad:-pad, pad:-pad].reshape(bs, -1)

        C = torch.einsum("bip,bjp->bij", A, A) / count + torch.eye(3, device=mr.device).unsqueeze(0) * alpha
        Atb = torch.einsum("bip,bp->bi", A, b) / count
        coeff = torch.linalg.solve(C, Atb)
        var = torch.mean(b**2, dim=1) - torch.mean(b, dim=1)**2
        dist = torch.mean(b**2, dim=1) + torch.einsum("bi,bj,bij->b", coeff, coeff, C) - 2*torch.einsum("bi,bi->b", coeff, Atb)
        sym = (var - dist)/var.clamp_min(beta)
        
        return sym.clamp(0, 1)
//  [python]
def run(self, us, mr, radius=9, alpha=1e-3, beta=1e-2):
        us = us.squeeze(1)
        mr = mr.squeeze(1)

        bs = mr.size(0)
        pad = (mr.size(1) - (2*radius+1)) // 2
        count = (2*radius+1)**3

        self.f = self.f.to(mr.device)

        grad = torch.linalg.norm(F.conv3d(mr.unsqueeze(1), self.f, padding=1), dim=1)

        A = torch.ones(bs, 3, count, device=mr.device)
        A[:, 0] = mr[:, pad:-pad, pad:-pad, pad:-pad].reshape(bs, -1)
        A[:, 1] = grad[:, pad:-pad, pad:-pad, pad:-pad].reshape(bs, -1)
        b = us[:, pad:-pad, pad:-pad, pad:-pad].reshape(bs, -1)

        C = torch.einsum("bip,bjp->bij", A, A) / count + torch.eye(3, device=mr.device).unsqueeze(0) * alpha
        Atb = torch.einsum("bip,bp->bi", A, b) / count
        coeff = torch.linalg.solve(C, Atb)
        var = torch.mean(b**2, dim=1) - torch.mean(b, dim=1)**2
        dist = torch.mean(b**2, dim=1) + torch.einsum("bi,bj,bij->b", coeff, coeff, C) - 2*torch.einsum("bi,bi->b", coeff, Atb)
        sym = (var - dist)/var.clamp_min(beta)
        
        return sym.clamp(0, 1)
//  [python]
def translation_error(t, t_gt):
    return torch.norm(t - t_gt, dim=1)
//  [python]
def translation_error(t, t_gt):
    return torch.linalg.norm(t - t_gt, dim=1)
//  [python]
def torsion_angle_loss(pred_torsions=None, true_torsions=None,
                       pred_points=None,  true_points=None,
                       alt_true_points=None, alt_true_torsions=None,
                       coeff=2., norm_coeff=1e-2, angle_mask=None):
    if true_torsions is not None and true_points is None:
        true_points = angle_to_point_in_circum(true_torsions)
    if alt_true_torsions is not None and alt_true_points is None:
        alt_true_points = angle_to_point_in_circum(alt_true_torsions)

    # calc norm of angles
    norm = torch.norm(pred_points, dim=-1)
    angle_norm_loss = norm_coeff * (1-norm).abs()

    # do L2 on unit circle
    pred_points = pred_points / norm.unsqueeze(-1)
    torsion_loss = torch.pow(pred_points - true_points, 2).sum(dim=-1)

    if alt_true_points is not None:
        torsion_loss = torch.minimum(
            torsion_loss,
            torch.pow(pred_points - alt_true_points, 2).sum(dim=-1)
        )
    if coeff != 2.:
        torsion_loss *= coeff/2

    if angle_mask is None:
        angle_mask = torch.ones(*pred_points.shape[:-1], dtype=torch.bool)

    return (torsion_loss + angle_norm_loss)[angle_mask].mean()
//  [python]
def torsion_angle_loss(pred_torsions=None, true_torsions=None,
                       pred_points=None,  true_points=None,
                       alt_true_points=None, alt_true_torsions=None,
                       coeff=2., norm_coeff=1e-2, angle_mask=None):
    if true_torsions is not None and true_points is None:
        true_points = angle_to_point_in_circum(true_torsions)
    if alt_true_torsions is not None and alt_true_points is None:
        alt_true_points = angle_to_point_in_circum(alt_true_torsions)

    # calc norm of angles
    norm = torch.linalg.norm(pred_points, dim=-1)
    angle_norm_loss = norm_coeff * (1-norm).abs()

    # do L2 on unit circle
    pred_points = pred_points / norm.unsqueeze(-1)
    torsion_loss = torch.pow(pred_points - true_points, 2).sum(dim=-1)

    if alt_true_points is not None:
        torsion_loss = torch.minimum(
            torsion_loss,
            torch.pow(pred_points - alt_true_points, 2).sum(dim=-1)
        )
    if coeff != 2.:
        torsion_loss *= coeff/2

    if angle_mask is None:
        angle_mask = torch.ones(*pred_points.shape[:-1], dtype=torch.bool)

    return (torsion_loss + angle_norm_loss)[angle_mask].mean()
//  [python]
def rmse(pts, T, T_gt):
    pts_pred = pts @ T[:, :3, :3].transpose(1, 2) + T[:, :3, 3].unsqueeze(1)
    pts_gt = pts @ T_gt[:, :3, :3].transpose(1, 2) + T_gt[:, :3, 3].unsqueeze(1)
    return torch.norm(pts_pred - pts_gt, dim=2).mean(dim=1)
//  [python]
def rmse(pts, T, T_gt):
    pts_pred = pts @ T[:, :3, :3].transpose(1, 2) + T[:, :3, 3].unsqueeze(1)
    pts_gt = pts @ T_gt[:, :3, :3].transpose(1, 2) + T_gt[:, :3, 3].unsqueeze(1)
    return torch.linalg.norm(pts_pred - pts_gt, dim=2).mean(dim=1)
//  [python]
def init_meta(self):
        if self.solved_meta is not None:
            return self.solved_meta
        task = torch.randn(self.sf_dim)
        task = task / torch.norm(task)
        task = task.cpu().numpy()
        meta = OrderedDict()
        meta['task'] = task
        return meta
//  [python]
def init_meta(self):
        if self.solved_meta is not None:
            return self.solved_meta
        task = torch.randn(self.sf_dim)
        task = task / torch.linalg.norm(task)
        task = task.cpu().numpy()
        meta = OrderedDict()
        meta['task'] = task
        return meta
//  [python]
def visualize(self, i):
        init_r_err = torch.acos((self.T_gt[i, :3, :3].trace() - 1) / 2) * 180 / math.pi
        init_t_err = torch.norm(self.T_gt[i, :3, 3])
        eye = torch.eye(4).unsqueeze(0).to(self.T_gt.device)
        init_rmse = rmse(self.pts1[i:i+1], eye, self.T_gt[i:i+1])[0]
        pts1_trans = self.pts1[i] @ self.T_12[i, :3, :3].T + self.T_12[i, :3, 3]
        fig = visualize([self.pts1[i], self.gamma1[i], self.pi1[i], self.mu1[i], self.sigma1[i],
                         self.pts2[i], self.gamma2[i], self.pi2[i], self.mu2[i], self.sigma2[i],
                         pts1_trans, init_r_err, init_t_err, init_rmse,
                         self.r_err[i], self.t_err[i], self.rmse[i]])
        return fig
//  [python]
def visualize(self, i):
        init_r_err = torch.acos((self.T_gt[i, :3, :3].trace() - 1) / 2) * 180 / math.pi
        init_t_err = torch.linalg.norm(self.T_gt[i, :3, 3])
        eye = torch.eye(4).unsqueeze(0).to(self.T_gt.device)
        init_rmse = rmse(self.pts1[i:i+1], eye, self.T_gt[i:i+1])[0]
        pts1_trans = self.pts1[i] @ self.T_12[i, :3, :3].T + self.T_12[i, :3, 3]
        fig = visualize([self.pts1[i], self.gamma1[i], self.pi1[i], self.mu1[i], self.sigma1[i],
                         self.pts2[i], self.gamma2[i], self.pi2[i], self.mu2[i], self.sigma2[i],
                         pts1_trans, init_r_err, init_t_err, init_rmse,
                         self.r_err[i], self.t_err[i], self.rmse[i]])
        return fig
//  [python]
def project_action(self, u, x):
        if self.isD0:
            Px = x @ self.P
            g = 2 * Px @ self.B
            neg_h = self.alpha * (Px * x).sum(-1) + \
                    2 * torch.norm(Px @ self.G, dim=1) * torch.norm(x @ self.C.T, dim=1) + \
                    2 * (Px @ self.A * x).sum(-1)
            u = self.proj_layer(u, g, -neg_h)
        else:
            Px = x @ self.P
            const = torch.norm(x @ self.P @ self.G, dim=1)
            A = self.D.expand(x.shape[0], self.D.shape[0], self.D.shape[1])
            b = x @ self.C.T
            c = (1 / const).unsqueeze(1) * (-Px @ self.B)
            d = -((2 * Px @ self.A + self.alpha * Px) * x).sum(-1) / (2 * const)

            u = self.proj_layer(u, A, b, c, d)

        return u
//  [python]
def project_action(self, u, x):
        if self.isD0:
            Px = x @ self.P
            g = 2 * Px @ self.B
            neg_h = self.alpha * (Px * x).sum(-1) + \
                    2 * torch.linalg.norm(Px @ self.G, dim=1) * torch.linalg.norm(x @ self.C.T, dim=1) + \
                    2 * (Px @ self.A * x).sum(-1)
            u = self.proj_layer(u, g, -neg_h)
        else:
            Px = x @ self.P
            const = torch.linalg.norm(x @ self.P @ self.G, dim=1)
            A = self.D.expand(x.shape[0], self.D.shape[0], self.D.shape[1])
            b = x @ self.C.T
            c = (1 / const).unsqueeze(1) * (-Px @ self.B)
            d = -((2 * Px @ self.A + self.alpha * Px) * x).sum(-1) / (2 * const)

            u = self.proj_layer(u, A, b, c, d)

        return u
//  [python]
def match_features(feat_src, feat_ref, metric='l2'):
	""" Compute pairwise distance between features

	Args:
		feat_src: (B, J, C)
		feat_ref: (B, K, C)
		metric: either 'angle' or 'l2' (squared euclidean)

	Returns:
		Matching matrix (B, J, K). i'th row describes how well the i'th point
		 in the src agrees with every point in the ref.
	"""
	assert feat_src.shape[-1] == feat_ref.shape[-1]

	if metric == 'l2':
		dist_matrix = square_distance(feat_src, feat_ref)
	elif metric == 'angle':
		feat_src_norm = feat_src / (torch.norm(feat_src, dim=-1, keepdim=True) + _EPS)
		feat_ref_norm = feat_ref / (torch.norm(feat_ref, dim=-1, keepdim=True) + _EPS)

		dist_matrix = angle_difference(feat_src_norm, feat_ref_norm)
	else:
		raise NotImplementedError

	return dist_matrix
//  [python]
def match_features(feat_src, feat_ref, metric='l2'):
	""" Compute pairwise distance between features

	Args:
		feat_src: (B, J, C)
		feat_ref: (B, K, C)
		metric: either 'angle' or 'l2' (squared euclidean)

	Returns:
		Matching matrix (B, J, K). i'th row describes how well the i'th point
		 in the src agrees with every point in the ref.
	"""
	assert feat_src.shape[-1] == feat_ref.shape[-1]

	if metric == 'l2':
		dist_matrix = square_distance(feat_src, feat_ref)
	elif metric == 'angle':
		feat_src_norm = feat_src / (torch.linalg.norm(feat_src, dim=-1, keepdim=True) + _EPS)
		feat_ref_norm = feat_ref / (torch.linalg.norm(feat_ref, dim=-1, keepdim=True) + _EPS)

		dist_matrix = angle_difference(feat_src_norm, feat_ref_norm)
	else:
		raise NotImplementedError

	return dist_matrix
//  [python]
def track(args):
    results = {}  # keyed on (seed, bias, num_frames)
    for seed in args.seed:
        torch.manual_seed(seed)
        observations, states, sensor_bias = generate_data(
            max(args.num_frames), args.num_sensors
        )
        for bias, num_frames in itertools.product(args.bias, args.num_frames):
            print(
                "tracking with seed={}, bias={}, num_frames={}".format(
                    seed, bias, num_frames
                )
            )
            model = Model(args.num_sensors)
            optim = Adam(model.parameters(), lr=args.lr, betas=(0.5, 0.8))
            losses = []
            for i in range(args.num_epochs):
                optim.zero_grad()
                log_prob, posterior = model(observations[:num_frames], add_bias=bias)
                loss = -log_prob
                loss.backward()
                losses.append(loss.item())
                if i % 10 == 0:
                    print(loss.item())
                optim.step()

            # Collect evaluation metrics.
            final_state_true = states[num_frames - 1]
            assert final_state_true.shape == (4,)
            final_pos_true = final_state_true[:2]
            final_vel_true = final_state_true[2:]

            final_state_est = posterior.loc
            assert final_state_est.shape == (4,)
            final_pos_est = final_state_est[:2]
            final_vel_est = final_state_est[2:]
            final_pos_error = float(torch.norm(final_pos_true - final_pos_est))
            final_vel_error = float(torch.norm(final_vel_true - final_vel_est))
            print("final_pos_error = {}".format(final_pos_error))

            results[seed, bias, num_frames] = {
                "args": args,
                "observations": observations[:num_frames],
                "states": states[:num_frames],
                "sensor_bias": sensor_bias,
                "losses": losses,
                "bias_scale": float(model.log_bias_scale.exp()),
                "obs_noise": float(model.log_obs_noise.exp()),
                "trans_noise": float(model.log_trans_noise.exp()),
                "final_state_estimate": posterior,
                "final_pos_error": final_pos_error,
                "final_vel_error": final_vel_error,
            }
        if args.metrics_filename:
            print("saving output to: {}".format(args.metrics_filename))
            torch.save(results, args.metrics_filename)
    return results
//  [python]
def track(args):
    results = {}  # keyed on (seed, bias, num_frames)
    for seed in args.seed:
        torch.manual_seed(seed)
        observations, states, sensor_bias = generate_data(
            max(args.num_frames), args.num_sensors
        )
        for bias, num_frames in itertools.product(args.bias, args.num_frames):
            print(
                "tracking with seed={}, bias={}, num_frames={}".format(
                    seed, bias, num_frames
                )
            )
            model = Model(args.num_sensors)
            optim = Adam(model.parameters(), lr=args.lr, betas=(0.5, 0.8))
            losses = []
            for i in range(args.num_epochs):
                optim.zero_grad()
                log_prob, posterior = model(observations[:num_frames], add_bias=bias)
                loss = -log_prob
                loss.backward()
                losses.append(loss.item())
                if i % 10 == 0:
                    print(loss.item())
                optim.step()

            # Collect evaluation metrics.
            final_state_true = states[num_frames - 1]
            assert final_state_true.shape == (4,)
            final_pos_true = final_state_true[:2]
            final_vel_true = final_state_true[2:]

            final_state_est = posterior.loc
            assert final_state_est.shape == (4,)
            final_pos_est = final_state_est[:2]
            final_vel_est = final_state_est[2:]
            final_pos_error = float(torch.linalg.norm(final_pos_true - final_pos_est))
            final_vel_error = float(torch.linalg.norm(final_vel_true - final_vel_est))
            print("final_pos_error = {}".format(final_pos_error))

            results[seed, bias, num_frames] = {
                "args": args,
                "observations": observations[:num_frames],
                "states": states[:num_frames],
                "sensor_bias": sensor_bias,
                "losses": losses,
                "bias_scale": float(model.log_bias_scale.exp()),
                "obs_noise": float(model.log_obs_noise.exp()),
                "trans_noise": float(model.log_trans_noise.exp()),
                "final_state_estimate": posterior,
                "final_pos_error": final_pos_error,
                "final_vel_error": final_vel_error,
            }
        if args.metrics_filename:
            print("saving output to: {}".format(args.metrics_filename))
            torch.save(results, args.metrics_filename)
    return results
//  [python]
def forward(self, raw_feat, rectified_feat):
        """
        get the rank of input tensors.
        :param raw_feat: tensor [n, c, h, w]
        :param rectified_feat: tensor [n, c, h, w]
        """
        assert raw_feat.size() == rectified_feat.size()
        batch_size, _, h, w = raw_feat.size()
        # average alone the channel
        raw_feat = torch.mean(raw_feat, dim=1, keepdim=False)  # [n, h, w]
        rectified_feat = torch.mean(
            rectified_feat, dim=1, keepdim=False)  # [n, h, w]
        # l2 normalize the feat
        flatten_raw_feat = raw_feat.view(batch_size, -1)  # [n, h*w]
        flatten_rectified_feat = rectified_feat.view(
            batch_size, -1)  # [n, h*w]
        raw_feat_n = torch.norm(
            flatten_raw_feat,
            p=2,
            dim=1,
            keepdim=True)  # [n, 1]
        rectified_feat_n = torch.norm(
            flatten_rectified_feat, p=2, dim=1, keepdim=True)  # [n, 1]
        # get the normalized feat
        normed_raw_feat = flatten_raw_feat.div(raw_feat_n).view(-1, h, w)
        normed_rectified_feat = flatten_rectified_feat.div(
            rectified_feat_n).view(-1, h, w)
        # get the rank of normalized feat
        batch_s1, batch_s2 = [], []
        for i in range(batch_size):
            u1, s1, v1 = torch.svd(normed_raw_feat[i])
            u2, s2, v2 = torch.svd(normed_rectified_feat[i])
            batch_s1.append(s1)
            batch_s2.append(s2)
        batch_s1 = torch.stack(batch_s1, dim=0)  # [b, min(h, w)]
        batch_s2 = torch.stack(batch_s2, dim=0)  # [b, min(h, w)]
        # keep the singular value > the tol
        batch_s1 = torch.gt(batch_s1, self.tol)
        batch_s2 = torch.gt(batch_s2, self.tol)
        # get the rank after clamping
        rank1 = batch_s1.sum(dim=1, keepdim=False)  # [b]
        rank2 = batch_s2.sum(dim=1, keepdim=False)  # [b]
        # get the loss
        y = rank1.new()
        y.resize_as_(rank1)
        if self.reverse:
            y.fill_(-1)
        else:
            y.fill_(1)
        loss = F.margin_ranking_loss(
            rank1,
            rank2,
            y,
            margin=self.margin,
            size_average=False,
            reduce=False)
        # loss / length of each word
        if self.size_average:
            loss = torch.sum(loss) / batch_size
        elif self.reduce:
            loss = torch.sum(loss)
        else:
            pass
        return loss
//  [python]
def forward(self, raw_feat, rectified_feat):
        """
        get the rank of input tensors.
        :param raw_feat: tensor [n, c, h, w]
        :param rectified_feat: tensor [n, c, h, w]
        """
        assert raw_feat.size() == rectified_feat.size()
        batch_size, _, h, w = raw_feat.size()
        # average alone the channel
        raw_feat = torch.mean(raw_feat, dim=1, keepdim=False)  # [n, h, w]
        rectified_feat = torch.mean(
            rectified_feat, dim=1, keepdim=False)  # [n, h, w]
        # l2 normalize the feat
        flatten_raw_feat = raw_feat.view(batch_size, -1)  # [n, h*w]
        flatten_rectified_feat = rectified_feat.view(
            batch_size, -1)  # [n, h*w]
        raw_feat_n = torch.linalg.norm(
            flatten_raw_feat,
            ord=2,
            dim=1,
            keepdim=True)  # [n, 1]
        rectified_feat_n = torch.linalg.norm(
            flatten_rectified_feat, ord=2, dim=1, keepdim=True)  # [n, 1]
        # get the normalized feat
        normed_raw_feat = flatten_raw_feat.div(raw_feat_n).view(-1, h, w)
        normed_rectified_feat = flatten_rectified_feat.div(
            rectified_feat_n).view(-1, h, w)
        # get the rank of normalized feat
        batch_s1, batch_s2 = [], []
        for i in range(batch_size):
            u1, s1, v1 = torch.svd(normed_raw_feat[i])
            u2, s2, v2 = torch.svd(normed_rectified_feat[i])
            batch_s1.append(s1)
            batch_s2.append(s2)
        batch_s1 = torch.stack(batch_s1, dim=0)  # [b, min(h, w)]
        batch_s2 = torch.stack(batch_s2, dim=0)  # [b, min(h, w)]
        # keep the singular value > the tol
        batch_s1 = torch.gt(batch_s1, self.tol)
        batch_s2 = torch.gt(batch_s2, self.tol)
        # get the rank after clamping
        rank1 = batch_s1.sum(dim=1, keepdim=False)  # [b]
        rank2 = batch_s2.sum(dim=1, keepdim=False)  # [b]
        # get the loss
        y = rank1.new()
        y.resize_as_(rank1)
        if self.reverse:
            y.fill_(-1)
        else:
            y.fill_(1)
        loss = F.margin_ranking_loss(
            rank1,
            rank2,
            y,
            margin=self.margin,
            size_average=False,
            reduce=False)
        # loss / length of each word
        if self.size_average:
            loss = torch.sum(loss) / batch_size
        elif self.reduce:
            loss = torch.sum(loss)
        else:
            pass
        return loss
//  [python]
def compute_intr_reward(self, task, next_obs, step):
        # maxent reward
        with torch.no_grad():
            rep = self.aps(next_obs, norm=False)
        reward = self.pbe(rep)
        intr_ent_reward = reward.reshape(-1, 1)

        # successor feature reward
        rep = rep / torch.norm(rep, dim=1, keepdim=True)
        intr_sf_reward = torch.einsum("bi,bi->b", task, rep).reshape(-1, 1)

        return intr_ent_reward, intr_sf_reward
//  [python]
def compute_intr_reward(self, task, next_obs, step):
        # maxent reward
        with torch.no_grad():
            rep = self.aps(next_obs, norm=False)
        reward = self.pbe(rep)
        intr_ent_reward = reward.reshape(-1, 1)

        # successor feature reward
        rep = rep / torch.linalg.norm(rep, dim=1, keepdim=True)
        intr_sf_reward = torch.einsum("bi,bi->b", task, rep).reshape(-1, 1)

        return intr_ent_reward, intr_sf_reward
//  [python]
def contract(x):
    # x: [N, 3]
    return (2 - 1 / (torch.norm(x, dim=-1, keepdim=True))) * x / torch.norm(x, dim=-1, keepdim=True)
//  [python]
def contract(x):
    # x: [N, 3]
    return (2 - 1 / (torch.linalg.norm(x, dim=-1, keepdim=True))) * x / torch.linalg.norm(x, dim=-1, keepdim=True)
//  [python]
def regress_meta(self, replay_iter, step):
        obs, reward = [], []
        batch_size = 0
        while batch_size < self.lstsq_batch_size:
            batch = next(replay_iter)
            batch_obs, _, batch_reward, *_ = utils.to_torch(batch, self.device)
            obs.append(batch_obs)
            reward.append(batch_reward)
            batch_size += batch_obs.size(0)
        obs, reward = torch.cat(obs, 0), torch.cat(reward, 0)

        obs = self.aug_and_encode(obs)
        rep = self.aps(obs)
        task = torch.linalg.lstsq(reward, rep)[0][:rep.size(1), :][0]
        task = task / torch.norm(task)
        task = task.cpu().numpy()
        meta = OrderedDict()
        meta['task'] = task

        # save for evaluation
        self.solved_meta = meta
        return meta
//  [python]
def regress_meta(self, replay_iter, step):
        obs, reward = [], []
        batch_size = 0
        while batch_size < self.lstsq_batch_size:
            batch = next(replay_iter)
            batch_obs, _, batch_reward, *_ = utils.to_torch(batch, self.device)
            obs.append(batch_obs)
            reward.append(batch_reward)
            batch_size += batch_obs.size(0)
        obs, reward = torch.cat(obs, 0), torch.cat(reward, 0)

        obs = self.aug_and_encode(obs)
        rep = self.aps(obs)
        task = torch.linalg.lstsq(reward, rep)[0][:rep.size(1), :][0]
        task = task / torch.linalg.norm(task)
        task = task.cpu().numpy()
        meta = OrderedDict()
        meta['task'] = task

        # save for evaluation
        self.solved_meta = meta
        return meta
//  [python]
def test_forward_frobenius_norm():
    """test_forward_frobenius_norm"""
    torch.set_grad_enabled(False)
    input_shape = [1, 3, 10, 10]

    class FroNorm1(Module):
        def forward(self, *args):
            return torch.norm(args[0])

    class FroNorm2(Module):
        def forward(self, *args):
            return torch.norm(args[0], p="fro", dim=None, keepdim=True)

    class FroNorm3(Module):
        def forward(self, *args):
            return torch.norm(args[0], p="fro", dim=(1), keepdim=True)

    class FroNorm4(Module):
        def forward(self, *args):
            return torch.norm(args[0], dim=None, keepdim=False)

    input_data = torch.rand(input_shape).float()
    verify_model(FroNorm1().float().eval(), input_data=input_data)
    verify_model(FroNorm2().float().eval(), input_data=input_data)
    verify_model(FroNorm3().float().eval(), input_data=input_data)
    verify_model(FroNorm4().float().eval(), input_data=input_data)
//  [python]
def test_forward_frobenius_norm():
    """test_forward_frobenius_norm"""
    torch.set_grad_enabled(False)
    input_shape = [1, 3, 10, 10]

    class FroNorm1(Module):
        def forward(self, *args):
            return torch.linalg.norm(args[0])

    class FroNorm2(Module):
        def forward(self, *args):
            return torch.linalg.norm(args[0], ord="fro", keepdim=True)

    class FroNorm3(Module):
        def forward(self, *args):
            return torch.linalg.norm(args[0], ord="fro", dim=(1), keepdim=True)

    class FroNorm4(Module):
        def forward(self, *args):
            return torch.linalg.norm(args[0], keepdim=False)

    input_data = torch.rand(input_shape).float()
    verify_model(FroNorm1().float().eval(), input_data=input_data)
    verify_model(FroNorm2().float().eval(), input_data=input_data)
    verify_model(FroNorm3().float().eval(), input_data=input_data)
    verify_model(FroNorm4().float().eval(), input_data=input_data)
//  [python]
def line_search(update, x0, g0, g, nstep=0, on=True):
    """
    `update` is the propsoed direction of update.

    Code adapted from scipy.
    """
    tmp_s = [0]
    tmp_g0 = [g0]
    tmp_phi = [torch.norm(g0) ** 2]
    s_norm = torch.norm(x0) / torch.norm(update)

    def phi(s, store=True):
        if s == tmp_s[0]:
            return tmp_phi[0]  # If the step size is so small... just return something
        x_est = x0 + s * update
        g0_new = g(x_est)
        phi_new = _safe_norm(g0_new) ** 2
        if store:
            tmp_s[0] = s
            tmp_g0[0] = g0_new
            tmp_phi[0] = phi_new
        return phi_new

    if on:
        s, phi1, ite = scalar_search_armijo(phi, tmp_phi[0], -tmp_phi[0], amin=1e-2)
    if (not on) or s is None:
        s = 1.0
        ite = 0

    x_est = x0 + s * update
    if s == tmp_s[0]:
        g0_new = tmp_g0[0]
    else:
        g0_new = g(x_est)
    return x_est, g0_new, x_est - x0, g0_new - g0, ite
//  [python]
def line_search(update, x0, g0, g, nstep=0, on=True):
    """
    `update` is the propsoed direction of update.

    Code adapted from scipy.
    """
    tmp_s = [0]
    tmp_g0 = [g0]
    tmp_phi = [torch.linalg.norm(g0) ** 2]
    s_norm = torch.linalg.norm(x0) / torch.linalg.norm(update)

    def phi(s, store=True):
        if s == tmp_s[0]:
            return tmp_phi[0]  # If the step size is so small... just return something
        x_est = x0 + s * update
        g0_new = g(x_est)
        phi_new = _safe_norm(g0_new) ** 2
        if store:
            tmp_s[0] = s
            tmp_g0[0] = g0_new
            tmp_phi[0] = phi_new
        return phi_new

    if on:
        s, phi1, ite = scalar_search_armijo(phi, tmp_phi[0], -tmp_phi[0], amin=1e-2)
    if (not on) or s is None:
        s = 1.0
        ite = 0

    x_est = x0 + s * update
    if s == tmp_s[0]:
        g0_new = tmp_g0[0]
    else:
        g0_new = g(x_est)
    return x_est, g0_new, x_est - x0, g0_new - g0, ite
//  [python]
def parameterization(means, covs):
    '''
    means: [B, N, 3]
    covs: [B, N, 3, 3]
    '''

    B, N, _ = means.shape
    means = means.reshape([-1, 3])
    if len(covs.shape) == 4:
        covs = covs.reshape(-1, 3, 3)
    else:
        covs = covs.reshape(-1, 3)
    contr_mask = (torch.norm(means, dim=-1, keepdim=True) > 1).detach()
    with torch.no_grad():
        jac = vmap(jacrev(contract))(means)
        print('11', jac.shape, covs.shape)
    means = torch.where(contr_mask, contract(means), means)
    covs = torch.where(contr_mask.unsqueeze(-1).expand(jac.shape), jac, covs)
    return means.reshape([B, N, 3]), covs.reshape([B, N, 3, 3])
//  [python]
def parameterization(means, covs):
    '''
    means: [B, N, 3]
    covs: [B, N, 3, 3]
    '''

    B, N, _ = means.shape
    means = means.reshape([-1, 3])
    if len(covs.shape) == 4:
        covs = covs.reshape(-1, 3, 3)
    else:
        covs = covs.reshape(-1, 3)
    contr_mask = (torch.linalg.norm(means, dim=-1, keepdim=True) > 1).detach()
    with torch.no_grad():
        jac = vmap(jacrev(contract))(means)
        print('11', jac.shape, covs.shape)
    means = torch.where(contr_mask, contract(means), means)
    covs = torch.where(contr_mask.unsqueeze(-1).expand(jac.shape), jac, covs)
    return means.reshape([B, N, 3]), covs.reshape([B, N, 3, 3])
//  [python]
def whitenapply_torch(X, m, P, dimensions=None):
    if not dimensions:
        dimensions = P.size(0)

    X = torch.mm(P[:dimensions, :], X - m)
    X = X / (torch.norm(X, dim=0, keepdim=True) + 1e-6)

    return X
//  [python]
def whitenapply_torch(X, m, P, dimensions=None):
    if not dimensions:
        dimensions = P.size(0)

    X = torch.mm(P[:dimensions, :], X - m)
    X = X / (torch.linalg.norm(X, dim=0, keepdim=True) + 1e-6)

    return X