jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def check_if_mirror_extended(self, lqs):
        """Check whether the input is a mirror-extended sequence.

        If mirror-extended, the i-th (i=0, ..., t-1) frame is equal to the
        (t-1-i)-th frame.

        Args:
            lqs (tensor): Input low quality (LQ) sequence with
                shape (n, t, c, h, w).
        """

        if lqs.size(1) % 2 == 0:
            lqs_1, lqs_2 = torch.chunk(lqs, 2, dim=1)
            if torch.norm(lqs_1 - lqs_2.flip(1)) == 0:
                self.is_mirror_extended = True
//  [python]
def check_if_mirror_extended(self, lqs):
        """Check whether the input is a mirror-extended sequence.

        If mirror-extended, the i-th (i=0, ..., t-1) frame is equal to the
        (t-1-i)-th frame.

        Args:
            lqs (tensor): Input low quality (LQ) sequence with
                shape (n, t, c, h, w).
        """

        if lqs.size(1) % 2 == 0:
            lqs_1, lqs_2 = torch.chunk(lqs, 2, dim=1)
            if torch.linalg.norm(lqs_1 - lqs_2.flip(1)) == 0:
                self.is_mirror_extended = True
//  [python]
def weighted_8points(x_in, logits):
    mask = logits[:, 0, :, 0] 
    weights = logits[:, 1, :, 0] 

    mask = torch.sigmoid(mask)
    weights = torch.exp(weights) * mask
    weights = weights / (torch.sum(weights, dim=-1, keepdim=True) + 1e-5)

    x_shp = x_in.shape
    x_in = x_in.squeeze(1)

    xx = torch.reshape(x_in, (x_shp[0], x_shp[2], 4)).permute(0, 2, 1).contiguous()

    X = torch.stack([
        xx[:, 2] * xx[:, 0], xx[:, 2] * xx[:, 1], xx[:, 2],
        xx[:, 3] * xx[:, 0], xx[:, 3] * xx[:, 1], xx[:, 3],
        xx[:, 0], xx[:, 1], torch.ones_like(xx[:, 0])
    ], dim=1).permute(0, 2, 1).contiguous()
    wX = torch.reshape(weights, (x_shp[0], x_shp[2], 1)) * X
    XwX = torch.matmul(X.permute(0, 2, 1).contiguous(), wX)

    # Recover essential matrix from self-adjoing eigen

    v = batch_symeig(XwX)
    e_hat = torch.reshape(v[:, :, 0], (x_shp[0], 9))
    e_hat = e_hat / torch.norm(e_hat, dim=1, keepdim=True)
    return e_hat
//  [python]
def weighted_8points(x_in, logits):
    mask = logits[:, 0, :, 0] 
    weights = logits[:, 1, :, 0] 

    mask = torch.sigmoid(mask)
    weights = torch.exp(weights) * mask
    weights = weights / (torch.sum(weights, dim=-1, keepdim=True) + 1e-5)

    x_shp = x_in.shape
    x_in = x_in.squeeze(1)

    xx = torch.reshape(x_in, (x_shp[0], x_shp[2], 4)).permute(0, 2, 1).contiguous()

    X = torch.stack([
        xx[:, 2] * xx[:, 0], xx[:, 2] * xx[:, 1], xx[:, 2],
        xx[:, 3] * xx[:, 0], xx[:, 3] * xx[:, 1], xx[:, 3],
        xx[:, 0], xx[:, 1], torch.ones_like(xx[:, 0])
    ], dim=1).permute(0, 2, 1).contiguous()
    wX = torch.reshape(weights, (x_shp[0], x_shp[2], 1)) * X
    XwX = torch.matmul(X.permute(0, 2, 1).contiguous(), wX)

    # Recover essential matrix from self-adjoing eigen

    v = batch_symeig(XwX)
    e_hat = torch.reshape(v[:, :, 0], (x_shp[0], 9))
    e_hat = e_hat / torch.linalg.norm(e_hat, dim=1, keepdim=True)
    return e_hat
//  [python]
def compute_franka_reward(
    reset_buf, progress_buf, successes, consecutive_successes, actions, states, reward_settings, max_episode_length
):
    # type: (Tensor, Tensor, Tensor, Tensor, Tensor, Dict[str, Tensor], Dict[str, float], float) -> Tuple[Tensor, Tensor, Tensor, Tensor]

    # Compute per-env physical parameters
    target_height = states["cubeB_size"] + states["cubeA_size"] / 2.0
    cubeA_size = states["cubeA_size"]
    cubeB_size = states["cubeB_size"]

    # distance from hand to the cubeA
    d = torch.norm(states["cubeA_pos_relative"], dim=-1)
    d_lf = torch.norm(states["cubeA_pos"] - states["eef_lf_pos"], dim=-1)
    d_rf = torch.norm(states["cubeA_pos"] - states["eef_rf_pos"], dim=-1)
    dist_reward = 1 - torch.tanh(10.0 * (d + d_lf + d_rf) / 3)

    # reward for lifting cubeA
    cubeA_height = states["cubeA_pos"][:, 2] - reward_settings["table_height"]
    cubeA_lifted = (cubeA_height - cubeA_size) > 0.04
    lift_reward = cubeA_lifted

    # how closely aligned cubeA is to cubeB (only provided if cubeA is lifted)
    offset = torch.zeros_like(states["cubeA_to_cubeB_pos"])
    offset[:, 2] = (cubeA_size + cubeB_size) / 2
    d_ab = torch.norm(states["cubeA_to_cubeB_pos"] + offset, dim=-1)
    align_reward = (1 - torch.tanh(10.0 * d_ab)) * cubeA_lifted

    # Dist reward is maximum of dist and align reward
    dist_reward = torch.max(dist_reward, align_reward)

    # final reward for stacking successfully (only if cubeA is close to target height and corresponding location, and gripper is not grasping)
    cubeA_align_cubeB = (torch.norm(states["cubeA_to_cubeB_pos"][:, :2], dim=-1) < 0.02)
    cubeA_on_cubeB = torch.abs(cubeA_height - target_height) < 0.02
    gripper_away_from_cubeA = (d > 0.04)
    stack_reward = cubeA_align_cubeB & cubeA_on_cubeB & gripper_away_from_cubeA

    # Compose rewards

    # We either provide the stack reward or the align + dist reward
    rewards = torch.where(
        stack_reward,
        reward_settings["r_stack_scale"] * stack_reward,
        reward_settings["r_dist_scale"] * dist_reward + reward_settings["r_lift_scale"] * lift_reward + reward_settings[
            "r_align_scale"] * align_reward,
    )

    # Compute resets
    successes = torch.where(stack_reward > 0, torch.ones_like(successes), successes)
    reset_buf = torch.where((progress_buf >= max_episode_length - 1) | (stack_reward > 0), torch.ones_like(reset_buf), reset_buf)

    consecutive_successes = torch.where(reset_buf > 0, successes * reset_buf, consecutive_successes).mean()
    return rewards, reset_buf, successes, consecutive_successes
//  [python]
def compute_franka_reward(
    reset_buf, progress_buf, successes, consecutive_successes, actions, states, reward_settings, max_episode_length
):
    # type: (Tensor, Tensor, Tensor, Tensor, Tensor, Dict[str, Tensor], Dict[str, float], float) -> Tuple[Tensor, Tensor, Tensor, Tensor]

    # Compute per-env physical parameters
    target_height = states["cubeB_size"] + states["cubeA_size"] / 2.0
    cubeA_size = states["cubeA_size"]
    cubeB_size = states["cubeB_size"]

    # distance from hand to the cubeA
    d = torch.linalg.norm(states["cubeA_pos_relative"], dim=-1)
    d_lf = torch.linalg.norm(states["cubeA_pos"] - states["eef_lf_pos"], dim=-1)
    d_rf = torch.linalg.norm(states["cubeA_pos"] - states["eef_rf_pos"], dim=-1)
    dist_reward = 1 - torch.tanh(10.0 * (d + d_lf + d_rf) / 3)

    # reward for lifting cubeA
    cubeA_height = states["cubeA_pos"][:, 2] - reward_settings["table_height"]
    cubeA_lifted = (cubeA_height - cubeA_size) > 0.04
    lift_reward = cubeA_lifted

    # how closely aligned cubeA is to cubeB (only provided if cubeA is lifted)
    offset = torch.zeros_like(states["cubeA_to_cubeB_pos"])
    offset[:, 2] = (cubeA_size + cubeB_size) / 2
    d_ab = torch.linalg.norm(states["cubeA_to_cubeB_pos"] + offset, dim=-1)
    align_reward = (1 - torch.tanh(10.0 * d_ab)) * cubeA_lifted

    # Dist reward is maximum of dist and align reward
    dist_reward = torch.max(dist_reward, align_reward)

    # final reward for stacking successfully (only if cubeA is close to target height and corresponding location, and gripper is not grasping)
    cubeA_align_cubeB = (torch.linalg.norm(states["cubeA_to_cubeB_pos"][:, :2], dim=-1) < 0.02)
    cubeA_on_cubeB = torch.abs(cubeA_height - target_height) < 0.02
    gripper_away_from_cubeA = (d > 0.04)
    stack_reward = cubeA_align_cubeB & cubeA_on_cubeB & gripper_away_from_cubeA

    # Compose rewards

    # We either provide the stack reward or the align + dist reward
    rewards = torch.where(
        stack_reward,
        reward_settings["r_stack_scale"] * stack_reward,
        reward_settings["r_dist_scale"] * dist_reward + reward_settings["r_lift_scale"] * lift_reward + reward_settings[
            "r_align_scale"] * align_reward,
    )

    # Compute resets
    successes = torch.where(stack_reward > 0, torch.ones_like(successes), successes)
    reset_buf = torch.where((progress_buf >= max_episode_length - 1) | (stack_reward > 0), torch.ones_like(reset_buf), reset_buf)

    consecutive_successes = torch.where(reset_buf > 0, successes * reset_buf, consecutive_successes).mean()
    return rewards, reset_buf, successes, consecutive_successes
//  [python]
def cal_logit_vec(self, entity_embedding, query_embedding):
        distance = entity_embedding - query_embedding
        logit = self.gamma - torch.norm(distance, p=1, dim=-1)
        return logit
//  [python]
def cal_logit_vec(self, entity_embedding, query_embedding):
        distance = entity_embedding - query_embedding
        logit = self.gamma - torch.linalg.norm(distance, ord=1, dim=-1)
        return logit
//  [python]
def cal_logit_beta(self, entity_embedding, query_dist):
        alpha_embedding, beta_embedding = torch.chunk(entity_embedding, 2, dim=-1)
        entity_dist = torch.distributions.beta.Beta(alpha_embedding, beta_embedding)
        logit = self.gamma - torch.norm(torch.distributions.kl.kl_divergence(entity_dist, query_dist), p=1, dim=-1)
        return logit
//  [python]
def cal_logit_beta(self, entity_embedding, query_dist):
        alpha_embedding, beta_embedding = torch.chunk(entity_embedding, 2, dim=-1)
        entity_dist = torch.distributions.beta.Beta(alpha_embedding, beta_embedding)
        logit = self.gamma - torch.linalg.norm(torch.distributions.kl.kl_divergence(entity_dist, query_dist), p=1, dim=-1)
        return logit
//  [python]
def cal_logit_gaussian(self, entity_embedding, query_center_embedding, query_offset_embedding):
        delta = (entity_embedding-query_center_embedding)
        inv_cov = torch.matmul(torch.transpose(query_offset_embedding,-2,-1),query_offset_embedding)
        distance = torch.matmul(torch.matmul(delta,inv_cov),torch.transpose(delta,-2,-1))
        logit = torch.norm(distance,p=1,dim=-1)
        return logit
//  [python]
def cal_logit_gaussian(self, entity_embedding, query_center_embedding, query_offset_embedding):
        delta = (entity_embedding-query_center_embedding)
        inv_cov = torch.matmul(torch.transpose(query_offset_embedding,-2,-1),query_offset_embedding)
        distance = torch.matmul(torch.matmul(delta,inv_cov),torch.transpose(delta,-2,-1))
        logit = torch.linalg.norm(distance,ord=1,dim=-1)
        return logit
//  [python]
def step(self, closure=None) -> float:
        starting_loss = closure(backward=True)

        flat_grad = self._gather_flat_grad()

        state = self.state
        if len(state) == 0:
            state['trust_radius'] = torch.full([1],
                                               self.initial_trust_radius,
                                               dtype=flat_grad.dtype,
                                               device=flat_grad.device)
        trust_radius = state['trust_radius']

        if self.opt_method == 'cg':
            param_step, hit_boundary = self._solve_subproblem_cg(
                starting_loss, flat_grad, trust_radius)
        else:
            param_step, hit_boundary = self._solve_subproblem_krylov(
                starting_loss, flat_grad, trust_radius)

        self.param_step = param_step

        if torch.norm(param_step).item() <= self.gtol:
            return starting_loss

        improvement_ratio = self._improvement_ratio(
            param_step, starting_loss, flat_grad, closure)

        if improvement_ratio.item() < 0.25:
            trust_radius.mul_(0.25)
        else:
            if improvement_ratio.item() > 0.75 and hit_boundary:
                trust_radius.mul_(2).clamp_(0.0, self.max_trust_radius)

        if improvement_ratio.item() <= self.eta:
            # If the improvement is not sufficient, then undo the update
            start_idx = 0
            for param in self._params:
                num_els = param.numel()
                curr_upd = param_step[start_idx:start_idx + num_els]
                param.data.add_(-curr_upd.view_as(param))
                start_idx += num_els

        self.steps += 1
        return starting_loss
//  [python]
def step(self, closure=None) -> float:
        starting_loss = closure(backward=True)

        flat_grad = self._gather_flat_grad()

        state = self.state
        if len(state) == 0:
            state['trust_radius'] = torch.full([1],
                                               self.initial_trust_radius,
                                               dtype=flat_grad.dtype,
                                               device=flat_grad.device)
        trust_radius = state['trust_radius']

        if self.opt_method == 'cg':
            param_step, hit_boundary = self._solve_subproblem_cg(
                starting_loss, flat_grad, trust_radius)
        else:
            param_step, hit_boundary = self._solve_subproblem_krylov(
                starting_loss, flat_grad, trust_radius)

        self.param_step = param_step

        if torch.linalg.norm(param_step).item() <= self.gtol:
            return starting_loss

        improvement_ratio = self._improvement_ratio(
            param_step, starting_loss, flat_grad, closure)

        if improvement_ratio.item() < 0.25:
            trust_radius.mul_(0.25)
        else:
            if improvement_ratio.item() > 0.75 and hit_boundary:
                trust_radius.mul_(2).clamp_(0.0, self.max_trust_radius)

        if improvement_ratio.item() <= self.eta:
            # If the improvement is not sufficient, then undo the update
            start_idx = 0
            for param in self._params:
                num_els = param.numel()
                curr_upd = param_step[start_idx:start_idx + num_els]
                param.data.add_(-curr_upd.view_as(param))
                start_idx += num_els

        self.steps += 1
        return starting_loss
//  [python]
def _root_finder(self, trust_radius, T_x, lanczos_g, loss, device):

        n_iter_nu, n_iter_r = 0, 0
        lambda_k, lambda_n, u_n = self._lambda_one_plus(T_x, device)
        lambda_const = self.lambda_const(lambda_k).to(device=device)
        if lambda_k == 0:  # T_x is positive definite
            _lambda = torch.tensor(
                [0], dtype=torch.float32, device=device)  # + lambda_const
        else:
            _lambda = lambda_k + lambda_const

        s, L = self._compute_s(_lambda=_lambda, lambda_const=lambda_const,
                               lanczos_g=lanczos_g, T_x=T_x, device=device)

        if norm(s) <= trust_radius:

            if _lambda == 0 or norm(s) == trust_radius:
                return s
            else:
                ta, tb = self.calc_boundaries(
                    iterate=s, direction=u_n, trust_radius=trust_radius)
                pa = s + ta * u_n
                pb = s + tb * u_n

                # Calculate the point on the boundary with the smallest value
                bound1_val = self._quad_model_krylov(lanczos_g, loss, pa, T_x)
                bound2_val = self._quad_model_krylov(lanczos_g, loss, pb, T_x)

                if bound1_val < bound2_val:
                    return pa
                else:
                    return pb

        while True:
            if self._converged(s, trust_radius) or norm(s) < torch.finfo(float).eps:
                break

            # w = torch.triangular_solve(
            #     s, L.T.to(device=device), upper=False).solution
            w = torch.linalg.solve_triangular(
                L.T.to(device=device), s, upper=False)
            _lambda = self._nu_next(_lambda, trust_radius, s, w)

            s, L = self._compute_s(_lambda, lambda_const,
                                   lanczos_g, T_x, device)

            n_iter_nu += 1
            if n_iter_nu > self.max_newton_iter - 1:  # self.max_krylov_dim:
                print(RuntimeWarning(
                    'Maximum number of newton iterations exceeded for _lambda: {}'.format(_lambda)))
                break

        return s
//  [python]
def _root_finder(self, trust_radius, T_x, lanczos_g, loss, device):

        n_iter_nu, n_iter_r = 0, 0
        lambda_k, lambda_n, u_n = self._lambda_one_plus(T_x, device)
        lambda_const = self.lambda_const(lambda_k).to(device=device)
        if lambda_k == 0:  # T_x is positive definite
            _lambda = torch.tensor(
                [0], dtype=torch.float32, device=device)  # + lambda_const
        else:
            _lambda = lambda_k + lambda_const

        s, L = self._compute_s(_lambda=_lambda, lambda_const=lambda_const,
                               lanczos_g=lanczos_g, T_x=T_x, device=device)

        if torch.linalg.norm(s) <= trust_radius:

            if _lambda == 0 or torch.linalg.norm(s) == trust_radius:
                return s
            else:
                ta, tb = self.calc_boundaries(
                    iterate=s, direction=u_n, trust_radius=trust_radius)
                pa = s + ta * u_n
                pb = s + tb * u_n

                # Calculate the point on the boundary with the smallest value
                bound1_val = self._quad_model_krylov(lanczos_g, loss, pa, T_x)
                bound2_val = self._quad_model_krylov(lanczos_g, loss, pb, T_x)

                if bound1_val < bound2_val:
                    return pa
                else:
                    return pb

        while True:
            if self._converged(s, trust_radius) or torch.linalg.norm(s) < torch.finfo(float).eps:
                break

            # w = torch.triangular_solve(
            #     s, L.T.to(device=device), upper=False).solution
            w = torch.linalg.solve_triangular(
                L.T.to(device=device), s, upper=False)
            _lambda = self._nu_next(_lambda, trust_radius, s, w)

            s, L = self._compute_s(_lambda, lambda_const,
                                   lanczos_g, T_x, device)

            n_iter_nu += 1
            if n_iter_nu > self.max_newton_iter - 1:  # self.max_krylov_dim:
                print(RuntimeWarning(
                    'Maximum number of newton iterations exceeded for _lambda: {}'.format(_lambda)))
                break

        return s
//  [python]
def get_tri_center_normal(tris: torch.Tensor) -> torch.Tensor:
    """
    Returns the center and normal of triangles
    Args:
        tris: tensor of shape (T, 3, 3)
    Returns:
        center: tensor of shape (T, 3)
        normal: tensor of shape (T, 3)
    """
    add_dim0 = False
    if tris.ndim == 2:
        tris = tris.unsqueeze(0)
        add_dim0 = True

    ctr = tris.mean(1)  # (T, 3)
    normals = torch.zeros_like(ctr)

    v0, v1, v2 = tris.unbind(1)  # 3 x (T, 3)

    # unvectorized solution
    T = tris.shape[0]
    for t in range(T):
        ns = torch.zeros((3, 3), device=tris.device)
        ns[0] = torch.cross(v0[t] - ctr[t], v1[t] - ctr[t], dim=-1)
        ns[1] = torch.cross(v0[t] - ctr[t], v2[t] - ctr[t], dim=-1)
        ns[2] = torch.cross(v1[t] - ctr[t], v2[t] - ctr[t], dim=-1)

        i = torch.norm(ns, dim=-1).argmax()
        normals[t] = ns[i]

    if add_dim0:
        ctr = ctr[0]
        normals = normals[0]
    normals = F.normalize(normals, dim=-1)
    return ctr, normals
//  [python]
def get_tri_center_normal(tris: torch.Tensor) -> torch.Tensor:
    """
    Returns the center and normal of triangles
    Args:
        tris: tensor of shape (T, 3, 3)
    Returns:
        center: tensor of shape (T, 3)
        normal: tensor of shape (T, 3)
    """
    add_dim0 = False
    if tris.ndim == 2:
        tris = tris.unsqueeze(0)
        add_dim0 = True

    ctr = tris.mean(1)  # (T, 3)
    normals = torch.zeros_like(ctr)

    v0, v1, v2 = tris.unbind(1)  # 3 x (T, 3)

    # unvectorized solution
    T = tris.shape[0]
    for t in range(T):
        ns = torch.zeros((3, 3), device=tris.device)
        ns[0] = torch.cross(v0[t] - ctr[t], v1[t] - ctr[t], dim=-1)
        ns[1] = torch.cross(v0[t] - ctr[t], v2[t] - ctr[t], dim=-1)
        ns[2] = torch.cross(v1[t] - ctr[t], v2[t] - ctr[t], dim=-1)

        i = torch.linalg.norm(ns, dim=-1).argmax()
        normals[t] = ns[i]

    if add_dim0:
        ctr = ctr[0]
        normals = normals[0]
    normals = F.normalize(normals, dim=-1)
    return ctr, normals
//  [python]
def coplanar_tri_plane(
    tri: torch.Tensor, plane: torch.Tensor, n: torch.Tensor, eps: float = DOT_EPS
):
    """
    Determines whether two triangle faces in 3D are coplanar
    Args:
            tri: tensor of shape (3, 3) of the vertices of the triangle
            plane: tensor of shape (4, 3) of the vertices of the plane
            n: tensor of shape (3,) of the unit "inside" direction on the plane
    Returns:
            is_coplanar: bool
    """
    tri_ctr, tri_n = get_tri_center_normal(tri)

    check1 = tri_n.dot(n).abs() > 1 - eps  # checks if parallel

    dist12 = torch.norm(tri.unsqueeze(1) - plane.unsqueeze(0), dim=-1)
    dist12_argmax = dist12.argmax()
    i1 = dist12_argmax // 4
    i2 = dist12_argmax % 4
    assert dist12[i1, i2] == dist12.max()

    check2 = F.normalize(tri[i1] - plane[i2], dim=0).dot(n).abs() < eps

    return check1 and check2
//  [python]
def coplanar_tri_plane(
    tri: torch.Tensor, plane: torch.Tensor, n: torch.Tensor, eps: float = DOT_EPS
):
    """
    Determines whether two triangle faces in 3D are coplanar
    Args:
            tri: tensor of shape (3, 3) of the vertices of the triangle
            plane: tensor of shape (4, 3) of the vertices of the plane
            n: tensor of shape (3,) of the unit "inside" direction on the plane
    Returns:
            is_coplanar: bool
    """
    tri_ctr, tri_n = get_tri_center_normal(tri)

    check1 = tri_n.dot(n).abs() > 1 - eps  # checks if parallel

    dist12 = torch.linalg.norm(tri.unsqueeze(1) - plane.unsqueeze(0), dim=-1)
    dist12_argmax = dist12.argmax()
    i1 = dist12_argmax // 4
    i2 = dist12_argmax % 4
    assert dist12[i1, i2] == dist12.max()

    check2 = F.normalize(tri[i1] - plane[i2], dim=0).dot(n).abs() < eps

    return check1 and check2
//  [python]
def coplanar_tri_faces(tri1: torch.Tensor, tri2: torch.Tensor, eps: float = DOT_EPS):
    """
    Determines whether two triangle faces in 3D are coplanar
    Args:
        tri1: tensor of shape (3, 3) of the vertices of the 1st triangle
        tri2: tensor of shape (3, 3) of the vertices of the 2nd triangle
    Returns:
        is_coplanar: bool
    """
    tri1_ctr, tri1_n = get_tri_center_normal(tri1)
    tri2_ctr, tri2_n = get_tri_center_normal(tri2)

    check1 = tri1_n.dot(tri2_n).abs() > 1 - eps  # checks if parallel

    dist12 = torch.norm(tri1.unsqueeze(1) - tri2.unsqueeze(0), dim=-1)
    dist12_argmax = dist12.argmax()
    i1 = dist12_argmax // 3
    i2 = dist12_argmax % 3
    assert dist12[i1, i2] == dist12.max()

    check2 = (
        F.normalize(tri1[i1] - tri2[i2], dim=0).dot(tri1_n).abs() < eps
    ) or F.normalize(tri1[i1] - tri2[i2], dim=0).dot(tri2_n).abs() < eps

    return check1 and check2
//  [python]
def coplanar_tri_faces(tri1: torch.Tensor, tri2: torch.Tensor, eps: float = DOT_EPS):
    """
    Determines whether two triangle faces in 3D are coplanar
    Args:
        tri1: tensor of shape (3, 3) of the vertices of the 1st triangle
        tri2: tensor of shape (3, 3) of the vertices of the 2nd triangle
    Returns:
        is_coplanar: bool
    """
    tri1_ctr, tri1_n = get_tri_center_normal(tri1)
    tri2_ctr, tri2_n = get_tri_center_normal(tri2)

    check1 = tri1_n.dot(tri2_n).abs() > 1 - eps  # checks if parallel

    dist12 = torch.linalg.norm(tri1.unsqueeze(1) - tri2.unsqueeze(0), dim=-1)
    dist12_argmax = dist12.argmax()
    i1 = dist12_argmax // 3
    i2 = dist12_argmax % 3
    assert dist12[i1, i2] == dist12.max()

    check2 = (
        F.normalize(tri1[i1] - tri2[i2], dim=0).dot(tri1_n).abs() < eps
    ) or F.normalize(tri1[i1] - tri2[i2], dim=0).dot(tri2_n).abs() < eps

    return check1 and check2
//  [python]
def rot6d_to_rotmat(rot6d, eps: float = 1e-8):
    a1 = rot6d[..., :3].clone()
    a2 = rot6d[..., 3:].clone()
    ind = torch.norm(a1, dim=-1) < eps
    a1[ind] = torch.tensor([1.0, 0.0, 0.0], device=a1.device)
    b1 = normalize(a1)

    b2 = normalize(a2 - (b1 * a2).sum(dim=-1).unsqueeze(-1) * b1)
    ind = torch.norm(b2, dim=-1) < eps
    b2[ind] = torch.tensor([0.0, 1.0, 0.0], device=b2.device)

    b3 = torch.cross(b1, b2, dim=-1)
    mat = torch.stack([b1, b2, b3], dim=-1)
    return mat
//  [python]
def rot6d_to_rotmat(rot6d, eps: float = 1e-8):
    a1 = rot6d[..., :3].clone()
    a2 = rot6d[..., 3:].clone()
    ind = torch.linalg.norm(a1, dim=-1) < eps
    a1[ind] = torch.tensor([1.0, 0.0, 0.0], device=a1.device)
    b1 = normalize(a1)

    b2 = normalize(a2 - (b1 * a2).sum(dim=-1).unsqueeze(-1) * b1)
    ind = torch.linalg.norm(b2, dim=-1) < eps
    b2[ind] = torch.tensor([0.0, 1.0, 0.0], device=b2.device)

    b3 = torch.cross(b1, b2, dim=-1)
    mat = torch.stack([b1, b2, b3], dim=-1)
    return mat
//  [python]
def test_random_matrix(features):
    torch.random.manual_seed(0)

    DRAWS = 100
    DIM = 10
    for _ in range(DRAWS):
        q = features._get_random_ortho_matrix(
            1, DIM, device=_device, norm_distribution=NormDistribution.Xi
        ).squeeze(0)

        # Check that the matrix is indeed orthonormal
        torch.allclose(
            torch.diag(q @ q.transpose(0, 1)),
            torch.diag(torch.ones(10, device=_device)),
        )

        # Check that the row norm is in the right ballpark (sqrt(dim))
        assert abs(torch.mean(torch.norm(q, dim=1)).item() - math.sqrt(DIM)) < 1.0
//  [python]
def test_random_matrix(features):
    torch.random.manual_seed(0)

    DRAWS = 100
    DIM = 10
    for _ in range(DRAWS):
        q = features._get_random_ortho_matrix(
            1, DIM, device=_device, norm_distribution=NormDistribution.Xi
        ).squeeze(0)

        # Check that the matrix is indeed orthonormal
        torch.allclose(
            torch.diag(q @ q.transpose(0, 1)),
            torch.diag(torch.ones(10, device=_device)),
        )

        # Check that the row norm is in the right ballpark (sqrt(dim))
        assert abs(torch.mean(torch.linalg.norm(q, dim=1)).item() - math.sqrt(DIM)) < 1.0
//  [python]
def get_plane_center_normal(planes: torch.Tensor) -> torch.Tensor:
    """
    Returns the center and normal of planes
    Args:
        planes: tensor of shape (P, 4, 3)
    Returns:
        center: tensor of shape (P, 3)
        normal: tensor of shape (P, 3)
    """
    add_dim0 = False
    if planes.ndim == 2:
        planes = planes.unsqueeze(0)
        add_dim0 = True

    ctr = planes.mean(1)  # (P, 3)
    normals = torch.zeros_like(ctr)

    v0, v1, v2, v3 = planes.unbind(1)  # 4 x (P, 3)

    # unvectorized solution
    P = planes.shape[0]
    for t in range(P):
        ns = torch.zeros((6, 3), device=planes.device)
        ns[0] = torch.cross(v0[t] - ctr[t], v1[t] - ctr[t], dim=-1)
        ns[1] = torch.cross(v0[t] - ctr[t], v2[t] - ctr[t], dim=-1)
        ns[2] = torch.cross(v0[t] - ctr[t], v3[t] - ctr[t], dim=-1)
        ns[3] = torch.cross(v1[t] - ctr[t], v2[t] - ctr[t], dim=-1)
        ns[4] = torch.cross(v1[t] - ctr[t], v3[t] - ctr[t], dim=-1)
        ns[5] = torch.cross(v2[t] - ctr[t], v3[t] - ctr[t], dim=-1)

        i = torch.norm(ns, dim=-1).argmax()
        normals[t] = ns[i]

    if add_dim0:
        ctr = ctr[0]
        normals = normals[0]
    normals = F.normalize(normals, dim=-1)
    return ctr, normals
//  [python]
def get_plane_center_normal(planes: torch.Tensor) -> torch.Tensor:
    """
    Returns the center and normal of planes
    Args:
        planes: tensor of shape (P, 4, 3)
    Returns:
        center: tensor of shape (P, 3)
        normal: tensor of shape (P, 3)
    """
    add_dim0 = False
    if planes.ndim == 2:
        planes = planes.unsqueeze(0)
        add_dim0 = True

    ctr = planes.mean(1)  # (P, 3)
    normals = torch.zeros_like(ctr)

    v0, v1, v2, v3 = planes.unbind(1)  # 4 x (P, 3)

    # unvectorized solution
    P = planes.shape[0]
    for t in range(P):
        ns = torch.zeros((6, 3), device=planes.device)
        ns[0] = torch.cross(v0[t] - ctr[t], v1[t] - ctr[t], dim=-1)
        ns[1] = torch.cross(v0[t] - ctr[t], v2[t] - ctr[t], dim=-1)
        ns[2] = torch.cross(v0[t] - ctr[t], v3[t] - ctr[t], dim=-1)
        ns[3] = torch.cross(v1[t] - ctr[t], v2[t] - ctr[t], dim=-1)
        ns[4] = torch.cross(v1[t] - ctr[t], v3[t] - ctr[t], dim=-1)
        ns[5] = torch.cross(v2[t] - ctr[t], v3[t] - ctr[t], dim=-1)

        i = torch.linalg.norm(ns, dim=-1).argmax()
        normals[t] = ns[i]

    if add_dim0:
        ctr = ctr[0]
        normals = normals[0]
    normals = F.normalize(normals, dim=-1)
    return ctr, normals
//  [python]
def compute_franka_reward(
    reset_buf, progress_buf, actions, states, reward_settings, max_episode_length
):
    # type: (Tensor, Tensor, Tensor, Dict[str, Tensor], Dict[str, float], float) -> Tuple[Tensor, Tensor]

    # Compute per-env physical parameters
    target_height = states["cubeB_size"] + states["cubeA_size"] / 2.0
    cubeA_size = states["cubeA_size"]
    cubeB_size = states["cubeB_size"]

    # distance from hand to the cubeA
    d = torch.norm(states["cubeA_pos_relative"], dim=-1)
    d_lf = torch.norm(states["cubeA_pos"] - states["eef_lf_pos"], dim=-1)
    d_rf = torch.norm(states["cubeA_pos"] - states["eef_rf_pos"], dim=-1)
    dist_reward = 1 - torch.tanh(10.0 * (d + d_lf + d_rf) / 3)

    # reward for lifting cubeA
    cubeA_height = states["cubeA_pos"][:, 2] - reward_settings["table_height"]
    cubeA_lifted = (cubeA_height - cubeA_size) > 0.04
    lift_reward = cubeA_lifted

    # how closely aligned cubeA is to cubeB (only provided if cubeA is lifted)
    offset = torch.zeros_like(states["cubeA_to_cubeB_pos"])
    offset[:, 2] = (cubeA_size + cubeB_size) / 2
    d_ab = torch.norm(states["cubeA_to_cubeB_pos"] + offset, dim=-1)
    align_reward = (1 - torch.tanh(10.0 * d_ab)) * cubeA_lifted

    # Dist reward is maximum of dist and align reward
    dist_reward = torch.max(dist_reward, align_reward)

    # final reward for stacking successfully (only if cubeA is close to target height and corresponding location, and gripper is not grasping)
    cubeA_align_cubeB = (torch.norm(states["cubeA_to_cubeB_pos"][:, :2], dim=-1) < 0.02)
    cubeA_on_cubeB = torch.abs(cubeA_height - target_height) < 0.02
    gripper_away_from_cubeA = (d > 0.04)
    stack_reward = cubeA_align_cubeB & cubeA_on_cubeB & gripper_away_from_cubeA

    # Compose rewards

    # We either provide the stack reward or the align + dist reward
    rewards = torch.where(
        stack_reward,
        reward_settings["r_stack_scale"] * stack_reward,
        reward_settings["r_dist_scale"] * dist_reward + reward_settings["r_lift_scale"] * lift_reward + reward_settings[
            "r_align_scale"] * align_reward,
    )

    # Compute resets
    reset_buf = torch.where((progress_buf >= max_episode_length - 1) | (stack_reward > 0), torch.ones_like(reset_buf), reset_buf)

    return rewards, reset_buf
//  [python]
def compute_franka_reward(
    reset_buf, progress_buf, actions, states, reward_settings, max_episode_length
):
    # type: (Tensor, Tensor, Tensor, Dict[str, Tensor], Dict[str, float], float) -> Tuple[Tensor, Tensor]

    # Compute per-env physical parameters
    target_height = states["cubeB_size"] + states["cubeA_size"] / 2.0
    cubeA_size = states["cubeA_size"]
    cubeB_size = states["cubeB_size"]

    # distance from hand to the cubeA
    d = torch.linalg.norm(states["cubeA_pos_relative"], dim=-1)
    d_lf = torch.linalg.norm(states["cubeA_pos"] - states["eef_lf_pos"], dim=-1)
    d_rf = torch.linalg.norm(states["cubeA_pos"] - states["eef_rf_pos"], dim=-1)
    dist_reward = 1 - torch.tanh(10.0 * (d + d_lf + d_rf) / 3)

    # reward for lifting cubeA
    cubeA_height = states["cubeA_pos"][:, 2] - reward_settings["table_height"]
    cubeA_lifted = (cubeA_height - cubeA_size) > 0.04
    lift_reward = cubeA_lifted

    # how closely aligned cubeA is to cubeB (only provided if cubeA is lifted)
    offset = torch.zeros_like(states["cubeA_to_cubeB_pos"])
    offset[:, 2] = (cubeA_size + cubeB_size) / 2
    d_ab = torch.linalg.norm(states["cubeA_to_cubeB_pos"] + offset, dim=-1)
    align_reward = (1 - torch.tanh(10.0 * d_ab)) * cubeA_lifted

    # Dist reward is maximum of dist and align reward
    dist_reward = torch.max(dist_reward, align_reward)

    # final reward for stacking successfully (only if cubeA is close to target height and corresponding location, and gripper is not grasping)
    cubeA_align_cubeB = (torch.linalg.norm(states["cubeA_to_cubeB_pos"][:, :2], dim=-1) < 0.02)
    cubeA_on_cubeB = torch.abs(cubeA_height - target_height) < 0.02
    gripper_away_from_cubeA = (d > 0.04)
    stack_reward = cubeA_align_cubeB & cubeA_on_cubeB & gripper_away_from_cubeA

    # Compose rewards

    # We either provide the stack reward or the align + dist reward
    rewards = torch.where(
        stack_reward,
        reward_settings["r_stack_scale"] * stack_reward,
        reward_settings["r_dist_scale"] * dist_reward + reward_settings["r_lift_scale"] * lift_reward + reward_settings[
            "r_align_scale"] * align_reward,
    )

    # Compute resets
    reset_buf = torch.where((progress_buf >= max_episode_length - 1) | (stack_reward > 0), torch.ones_like(reset_buf), reset_buf)

    return rewards, reset_buf
//  [python]
def backward(ctx, grad_output):
        M, M_sqrt_inv,  I = ctx.saved_tensors
        M_inv = M_sqrt_inv.bmm(M_sqrt_inv)
        grad_lya = - M_inv.bmm(grad_output).bmm(M_inv)
        norm_sqrt_inv = torch.norm(M_sqrt_inv)
        b = M_sqrt_inv / norm_sqrt_inv
        c = grad_lya / norm_sqrt_inv
        for i in range(8):
            #In case you might terminate the iteration by checking convergence
            #if th.norm(b-I)<1e-4:
            #    break
            b_2 = b.mm(b)
            c = 0.5 * (c.mm(3.0 * I - b_2) - b_2.mm(c) + b.mm(c).mm(b))
            b = 0.5 * b.mm(3.0 * I - b_2)
        grad_input = 0.5 * c
        return grad_input
//  [python]
def backward(ctx, grad_output):
        M, M_sqrt_inv,  I = ctx.saved_tensors
        M_inv = M_sqrt_inv.bmm(M_sqrt_inv)
        grad_lya = - M_inv.bmm(grad_output).bmm(M_inv)
        norm_sqrt_inv = torch.linalg.norm(M_sqrt_inv)
        b = M_sqrt_inv / norm_sqrt_inv
        c = grad_lya / norm_sqrt_inv
        for i in range(8):
            #In case you might terminate the iteration by checking convergence
            #if th.norm(b-I)<1e-4:
            #    break
            b_2 = b.mm(b)
            c = 0.5 * (c.mm(3.0 * I - b_2) - b_2.mm(c) + b.mm(c).mm(b))
            b = 0.5 * b.mm(3.0 * I - b_2)
        grad_input = 0.5 * c
        return grad_input
//  [python]
def backward(self, losses, **kwargs):
        self.update_weights_every = kwargs['update_weights_every']
        self.optim_niter = kwargs['optim_niter']
        self.max_norm = kwargs['max_norm']
        
        if self.step == 0:
            self._init_optim_problem()
        if (self.step % self.update_weights_every) == 0:
            self.step += 1

            if self.rep_grad:
                raise ValueError('No support method Nash_MTL with representation gradients (rep_grad=True)')
            else:
                self._compute_grad_dim()
                grads = self._compute_grad(losses, mode='autograd')
            
            GTG = torch.mm(grads, grads.t())
            self.normalization_factor = torch.norm(GTG).detach().cpu().numpy().reshape((1,))
            GTG = GTG / self.normalization_factor.item()
            alpha = self.solve_optimization(GTG.cpu().detach().numpy())
        else:
            self.step += 1
            alpha = self.prvs_alpha
            
        alpha = torch.from_numpy(alpha).to(torch.float32).to(self.device)
        torch.sum(alpha*losses).backward()

        if self.max_norm > 0:
            torch.nn.utils.clip_grad_norm_(self.get_share_params(), self.max_norm)
        
        return alpha.detach().cpu().numpy()
//  [python]
def backward(self, losses, **kwargs):
        self.update_weights_every = kwargs['update_weights_every']
        self.optim_niter = kwargs['optim_niter']
        self.max_norm = kwargs['max_norm']
        
        if self.step == 0:
            self._init_optim_problem()
        if (self.step % self.update_weights_every) == 0:
            self.step += 1

            if self.rep_grad:
                raise ValueError('No support method Nash_MTL with representation gradients (rep_grad=True)')
            else:
                self._compute_grad_dim()
                grads = self._compute_grad(losses, mode='autograd')
            
            GTG = torch.mm(grads, grads.t())
            self.normalization_factor = torch.linalg.norm(GTG).detach().cpu().numpy().reshape((1,))
            GTG = GTG / self.normalization_factor.item()
            alpha = self.solve_optimization(GTG.cpu().detach().numpy())
        else:
            self.step += 1
            alpha = self.prvs_alpha
            
        alpha = torch.from_numpy(alpha).to(torch.float32).to(self.device)
        torch.sum(alpha*losses).backward()

        if self.max_norm > 0:
            torch.nn.utils.clip_grad_norm_(self.get_share_params(), self.max_norm)
        
        return alpha.detach().cpu().numpy()
//  [python]
def setup(self, x0: torch.Tensor, y0: torch.Tensor, func: Callable):
        """Setup the Jacobian for the rootfinder.

        Parameters
        ----------
        x0
            The initial guess of the root.
        y0
            The function value at the initial guess.
        func
            The function to find the root.

        """
        self.x_prev = x0
        self.y_prev = y0

        if self.max_rank is None:
            self.max_rank = float('inf')

        if self.alpha is None:
            normy0 = torch.norm(y0)
            ones = torch.ones_like(normy0)
            if normy0:
                self.alpha = 0.5 * torch.max(torch.norm(x0), ones) / normy0
            else:
                self.alpha = ones

        if self.uv0 == "svd":
            self.uv0 = _get_svd_uv0(func, x0)

        # setup the approximate inverse Jacobian
        self.Gm = LowRankMatrix(-self.alpha, self.uv0, "restart")
//  [python]
def setup(self, x0: torch.Tensor, y0: torch.Tensor, func: Callable):
        """Setup the Jacobian for the rootfinder.

        Parameters
        ----------
        x0
            The initial guess of the root.
        y0
            The function value at the initial guess.
        func
            The function to find the root.

        """
        self.x_prev = x0
        self.y_prev = y0

        if self.max_rank is None:
            self.max_rank = float('inf')

        if self.alpha is None:
            normy0 = torch.linalg.norm(y0)
            ones = torch.ones_like(normy0)
            if normy0:
                self.alpha = 0.5 * torch.max(torch.linalg.norm(x0), ones) / normy0
            else:
                self.alpha = ones

        if self.uv0 == "svd":
            self.uv0 = _get_svd_uv0(func, x0)

        # setup the approximate inverse Jacobian
        self.Gm = LowRankMatrix(-self.alpha, self.uv0, "restart")
//  [python]
def norm_distribution(self, x, position, process):
    #print(copyx.shape)
    copyx = x.transpose(0,1) #B,T,C
    items = min(copyx.shape[0],2)
    for i in range(items):
        temx = copyx[i] #T,C
        len_x = self.mask[i].sum()
        len_x = len_x.int()
        temx = temx[:len_x] #len_x, C
        bag = torch.zeros(self.max_len)
        bag[:len_x] = torch.norm(temx,dim=1)
        bag[-1] = len_x.float()
        self.d_norm_distribution[process][position].append(bag.reshape([1,-1]))
//  [python]
def norm_distribution(self, x, position, process):
    #print(copyx.shape)
    copyx = x.transpose(0,1) #B,T,C
    items = min(copyx.shape[0],2)
    for i in range(items):
        temx = copyx[i] #T,C
        len_x = self.mask[i].sum()
        len_x = len_x.int()
        temx = temx[:len_x] #len_x, C
        bag = torch.zeros(self.max_len)
        bag[:len_x] = torch.linalg.norm(temx,dim=1)
        bag[-1] = len_x.float()
        self.d_norm_distribution[process][position].append(bag.reshape([1,-1]))
//  [python]
def get_rays(directions, c2w):
  """Get ray origin and normalized directions in world coordinate for all pixels in one image.

    Reference: https://www.scratchapixel.com/lessons/3d-basic-rendering/
               ray-tracing-generating-camera-rays/standard-coordinate-systems

    Inputs:
        directions: (H, W, 3) precomputed ray directions in camera coordinate
        c2w: (3, 4) transformation matrix from camera coordinate to world
        coordinate

    Outputs:
        rays_o: (H*W, 3), the origin of the rays in world coordinate
        rays_d: (H*W, 3), the normalized direction of the rays in world
        coordinate
  """
  # Rotate ray directions from camera coordinate to the world coordinate
  rays_d = directions @ c2w[:, :3].T  # (H, W, 3)
  rays_d /= torch.norm(rays_d, dim=-1, keepdim=True)
  # The origin of all rays is the camera origin in world coordinate
  rays_o = c2w[:, 3].expand(rays_d.shape)  # (H, W, 3)

  rays_d = rays_d.view(-1, 3)
  rays_o = rays_o.view(-1, 3)

  return rays_o, rays_d
//  [python]
def get_rays(directions, c2w):
  """Get ray origin and normalized directions in world coordinate for all pixels in one image.

    Reference: https://www.scratchapixel.com/lessons/3d-basic-rendering/
               ray-tracing-generating-camera-rays/standard-coordinate-systems

    Inputs:
        directions: (H, W, 3) precomputed ray directions in camera coordinate
        c2w: (3, 4) transformation matrix from camera coordinate to world
        coordinate

    Outputs:
        rays_o: (H*W, 3), the origin of the rays in world coordinate
        rays_d: (H*W, 3), the normalized direction of the rays in world
        coordinate
  """
  # Rotate ray directions from camera coordinate to the world coordinate
  rays_d = directions @ c2w[:, :3].T  # (H, W, 3)
  rays_d /= torch.linalg.norm(rays_d, dim=-1, keepdim=True)
  # The origin of all rays is the camera origin in world coordinate
  rays_o = c2w[:, 3].expand(rays_d.shape)  # (H, W, 3)

  rays_d = rays_d.view(-1, 3)
  rays_o = rays_o.view(-1, 3)

  return rays_o, rays_d
//  [python]
def frobenius_distance_torch(x1, x2, diagonal_distance=False):
    """
    Compute the Frobenius distance between matrix points

    Parameters
    ----------
    :param x1: set of matrices              (N1 x d1 x d2 or b1 x ... x bk x N1 x d1 x d2)
    :param x2: set of matrices              (N2 x d1 x d2 or b1 x ... x bk x N2 x d1 x d2)

    Optional parameters
    -------------------
    :param diagonal_distance: Whole distance matrix, or just the diagonal? If True, we must have `x1 == x2`.

    Returns
    -------
    :return: matrix of Frobenius distance between points in x1 and x2       (N1 x N2 or b1 x ... x bk x N1 x N2)
    """
    # If diag, x1 must be equal to x2 and we can return zeros.
    if diagonal_distance is True:
        shape = list(x2.shape)[:-2]
        shape.append(1)
        return torch.zeros(shape, dtype=x1.dtype)

    # Expand dimensions to compute all matrix-matrix distances
    x1 = x1.unsqueeze(-3)
    x2 = x2.unsqueeze(-4)

    # Repeat x and y data along 1- and 0- dimensions to have ndata_x x ndata_y x dim x dim arrays
    x1 = torch.cat(x2.shape[-3] * [x1], dim=-3)
    x2 = torch.cat(x1.shape[-4] * [x2], dim=-4)

    # Compute the distance between each pair of matrices
    return torch.norm(torch.add(x1, -x2) + 1e-15, dim=[-2, -1]).double()
//  [python]
def frobenius_distance_torch(x1, x2, diagonal_distance=False):
    """
    Compute the Frobenius distance between matrix points

    Parameters
    ----------
    :param x1: set of matrices              (N1 x d1 x d2 or b1 x ... x bk x N1 x d1 x d2)
    :param x2: set of matrices              (N2 x d1 x d2 or b1 x ... x bk x N2 x d1 x d2)

    Optional parameters
    -------------------
    :param diagonal_distance: Whole distance matrix, or just the diagonal? If True, we must have `x1 == x2`.

    Returns
    -------
    :return: matrix of Frobenius distance between points in x1 and x2       (N1 x N2 or b1 x ... x bk x N1 x N2)
    """
    # If diag, x1 must be equal to x2 and we can return zeros.
    if diagonal_distance is True:
        shape = list(x2.shape)[:-2]
        shape.append(1)
        return torch.zeros(shape, dtype=x1.dtype)

    # Expand dimensions to compute all matrix-matrix distances
    x1 = x1.unsqueeze(-3)
    x2 = x2.unsqueeze(-4)

    # Repeat x and y data along 1- and 0- dimensions to have ndata_x x ndata_y x dim x dim arrays
    x1 = torch.cat(x2.shape[-3] * [x1], dim=-3)
    x2 = torch.cat(x1.shape[-4] * [x2], dim=-4)

    # Compute the distance between each pair of matrices
    return torch.linalg.norm(torch.add(x1, -x2) + 1e-15, dim=[-2, -1]).double()
//  [python]
def contractive_loss(input, target, lambd, ae, reduction: str):
    """
    Actual function computing the loss of a contractive autoencoder
    :param input: (Tensor)
    :param target: (Tensor)
    :param lambd: (float) regularization parameter
    :param ae: (DeepAutoencoder) the model itself, used to get it's weights
    :param reduction: (str) type of reduction {'mean' | 'sum'}
    :raises: ValueError
    :return: the loss
    """
    term1 = (input - target) ** 2
    enc_weights = [ae.encoder[i].weight for i in reversed(range(1, len(ae.encoder), 2))]
    term2 = lambd * torch.norm(torch.chain_matmul(*enc_weights))
    contr_loss = torch.mean(term1 + term2, 0)
    if reduction == 'mean':
        return torch.mean(contr_loss)
    elif reduction == 'sum':
        return torch.sum(contr_loss)
    else:
        raise ValueError(f"value for 'reduction' must be 'mean' or 'sum', got {reduction}")
//  [python]
def contractive_loss(input, target, lambd, ae, reduction: str):
    """
    Actual function computing the loss of a contractive autoencoder
    :param input: (Tensor)
    :param target: (Tensor)
    :param lambd: (float) regularization parameter
    :param ae: (DeepAutoencoder) the model itself, used to get it's weights
    :param reduction: (str) type of reduction {'mean' | 'sum'}
    :raises: ValueError
    :return: the loss
    """
    term1 = (input - target) ** 2
    enc_weights = [ae.encoder[i].weight for i in reversed(range(1, len(ae.encoder), 2))]
    term2 = lambd * torch.linalg.norm(torch.chain_matmul(*enc_weights))
    contr_loss = torch.mean(term1 + term2, 0)
    if reduction == 'mean':
        return torch.mean(contr_loss)
    elif reduction == 'sum':
        return torch.sum(contr_loss)
    else:
        raise ValueError(f"value for 'reduction' must be 'mean' or 'sum', got {reduction}")
//  [python]
def spd_product_manifold_distance_torch(x1, x2, diagonal_distance=False):
    """
    Compute the distance between SPD points considered as a product of a Euclidean and a sphere manifold.

    Parameters
    ----------
    :param x1: set of SPD matrices      (N1 x d x d or b1 x ... x bk x N1 x d x d)
    :param x2: set of SPD matrices      (N2 x d x d or b1 x ... x bk x N1 x d x d)

    Optional parameters
    -------------------
    :param diagonal_distance: Whole distance matrix, or just the diagonal? If True, we must have `x1 == x2`.

    Returns
    -------
    :return: matrix of product-of-manifolds distance between SPD points in x1 and x2  (N1 x N2 or b1 x ... x bk x N1 x N2)
    """
    # If diag, x1 must be equal to x2 and we can return zeros.
    if diagonal_distance is True:
        shape = list(x2.shape)[:-2]
        shape.append(1)
        return torch.zeros(shape, dtype=x1.dtype)

    # Decompose SPD matrices in sphere (from SO(n)) and Euclidean parts using eigenvalue decompositions
    x_euclidean1, x_s1 = spd_to_sphere_and_euclidean_torch(x1)
    x_euclidean2, x_s2 = spd_to_sphere_and_euclidean_torch(x2)

    # Compute sphere distance
    distance_sphere = sphere_distance_torch(x_s1, x_s2)

    # Compute Euclidean distance
    distance_euclidean = torch.norm(x_euclidean1 - x_euclidean2, dim=-1).unsqueeze(-1)

    # Compute product-of-manifold distance
    distance = torch.sqrt(distance_sphere**2 + distance_euclidean**2)

    return distance.double()