jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def __init__(
        self,
        rot_mats: Optional[torch.Tensor] = None,
        quats: Optional[torch.Tensor] = None,
        normalize_quats: bool = True,
    ):
        if (rot_mats is None and quats is None) or (
            rot_mats is not None and quats is not None
        ):
            raise ValueError("Exactly one input argument must be specified")

        if (rot_mats is not None and rot_mats.shape[-2:] != (3, 3)) or (
            quats is not None and quats.shape[-1] != 4
        ):
            raise ValueError(
                "Incorrectly shaped rotation matrix or quaternion"
            )

        # Force full-precision
        if quats is not None:
            quats = quats.to(dtype=torch.float32)
        if rot_mats is not None:
            rot_mats = rot_mats.to(dtype=torch.float32)

        if quats is not None and normalize_quats:
            quats = quats / torch.linalg.norm(quats, dim=-1, keepdim=True)

        self._rot_mats = rot_mats
        self._quats = quats
//  [python]
def test_circle_grid_sampling_with_d():
    C = Circle(R2('x'), [0, 0], 4)
    points = C.sample_grid(d=1)
    assert all(t.norm(points.as_tensor, dim=1) <= 4.0)
//  [python]
def test_circle_grid_sampling_with_d():
    C = Circle(R2('x'), [0, 0], 4)
    points = C.sample_grid(d=1)
    assert all(torch.linalg.norm(points.as_tensor, dim=1) <= 4.0)
//  [python]
def test_circle_boundary_random_sampling_with_d():
    C = Circle(R2('x'), [0, 0], 4).boundary
    points = C.sample_random_uniform(d=15)
    assert all(torch.isclose(points.as_tensor.norm(dim=1), torch.tensor(4.0)))
//  [python]
def test_circle_boundary_random_sampling_with_d():
    C = Circle(R2('x'), [0, 0], 4).boundary
    points = C.sample_random_uniform(d=15)
    assert all(torch.isclose(torch.linalg.norm(points.as_tensor, dim=1), torch.tensor(4.0)))
//  [python]
def test_circle_boundary_random_sampling_with_n():
    C = Circle(R2('x'), [0, 0], 4).boundary
    points = C.sample_random_uniform(n=10)
    assert points.as_tensor.shape == (10, 2)
    assert all(torch.isclose(torch.norm(points.as_tensor, dim=1),
                             torch.tensor(4.0)))
//  [python]
def test_circle_boundary_random_sampling_with_n():
    C = Circle(R2('x'), [0, 0], 4).boundary
    points = C.sample_random_uniform(n=10)
    assert points.as_tensor.shape == (10, 2)
    assert all(torch.isclose(torch.linalg.norm(points.as_tensor, dim=1),
                             torch.tensor(4.0)))
//  [python]
def up_sample(self, rays_o, rays_d, z_vals, sdf, n_importance, inv_s):
        """
        Up sampling give a fixed inv_s
        """
        batch_size, n_samples = z_vals.shape
        pts = rays_o[:, None, :] + rays_d[:, None, :] * z_vals[..., :, None] # n_rays, n_samples, 3
        radius = torch.norm(pts, ord=2, dim=-1, keepdim=False)
        inside_sphere = (radius[:, :-1] < 1.0) | (radius[:, 1:] < 1.0)
        sdf = sdf.reshape(batch_size, n_samples)
        prev_sdf, next_sdf = sdf[:, :-1], sdf[:, 1:]
        prev_z_vals, next_z_vals = z_vals[:, :-1], z_vals[:, 1:]
        mid_sdf = (prev_sdf + next_sdf) * 0.5
        cos_val = (next_sdf - prev_sdf) / (next_z_vals - prev_z_vals + 1e-5)

        # ----------------------------------------------------------------------------------------------------------
        # Use min value of [ cos, prev_cos ]
        # Though it makes the sampling (not rendering) a little bit biased, this strategy can make the sampling more
        # robust when meeting situations like below:
        #
        # SDF
        # ^
        # |\          -----x----...
        # | \        /
        # |  x      x
        # |---\----/-------------> 0 level
        # |    \  /
        # |     \/
        # |
        # ----------------------------------------------------------------------------------------------------------
        prev_cos_val = torch.cat([torch.zeros([batch_size, 1]), cos_val[:, :-1]], dim=-1)
        cos_val = torch.stack([prev_cos_val, cos_val], dim=-1)
        cos_val, _ = torch.min(cos_val, dim=-1, keepdim=False)
        cos_val = cos_val.clip(-1e3, 0.0) * inside_sphere

        dist = (next_z_vals - prev_z_vals)
        prev_esti_sdf = mid_sdf - cos_val * dist * 0.5
        next_esti_sdf = mid_sdf + cos_val * dist * 0.5
        prev_cdf = torch.sigmoid(prev_esti_sdf * inv_s)
        next_cdf = torch.sigmoid(next_esti_sdf * inv_s)
        alpha = (prev_cdf - next_cdf + 1e-5) / (prev_cdf + 1e-5)
        weights = alpha * torch.cumprod(
            torch.cat([torch.ones([batch_size, 1]), 1. - alpha + 1e-7], -1), -1)[:, :-1]

        z_samples = sample_pdf(z_vals, weights, n_importance, det=True).detach()
        return z_samples
//  [python]
def up_sample(self, rays_o, rays_d, z_vals, sdf, n_importance, inv_s):
        """
        Up sampling give a fixed inv_s
        """
        batch_size, n_samples = z_vals.shape
        pts = rays_o[:, None, :] + rays_d[:, None, :] * z_vals[..., :, None] # n_rays, n_samples, 3
        radius = torch.linalg.norm(pts, ord=2, dim=-1, keepdim=False)
        inside_sphere = (radius[:, :-1] < 1.0) | (radius[:, 1:] < 1.0)
        sdf = sdf.reshape(batch_size, n_samples)
        prev_sdf, next_sdf = sdf[:, :-1], sdf[:, 1:]
        prev_z_vals, next_z_vals = z_vals[:, :-1], z_vals[:, 1:]
        mid_sdf = (prev_sdf + next_sdf) * 0.5
        cos_val = (next_sdf - prev_sdf) / (next_z_vals - prev_z_vals + 1e-5)

        # ----------------------------------------------------------------------------------------------------------
        # Use min value of [ cos, prev_cos ]
        # Though it makes the sampling (not rendering) a little bit biased, this strategy can make the sampling more
        # robust when meeting situations like below:
        #
        # SDF
        # ^
        # |\          -----x----...
        # | \        /
        # |  x      x
        # |---\----/-------------> 0 level
        # |    \  /
        # |     \/
        # |
        # ----------------------------------------------------------------------------------------------------------
        prev_cos_val = torch.cat([torch.zeros([batch_size, 1]), cos_val[:, :-1]], dim=-1)
        cos_val = torch.stack([prev_cos_val, cos_val], dim=-1)
        cos_val, _ = torch.min(cos_val, dim=-1, keepdim=False)
        cos_val = cos_val.clip(-1e3, 0.0) * inside_sphere

        dist = (next_z_vals - prev_z_vals)
        prev_esti_sdf = mid_sdf - cos_val * dist * 0.5
        next_esti_sdf = mid_sdf + cos_val * dist * 0.5
        prev_cdf = torch.sigmoid(prev_esti_sdf * inv_s)
        next_cdf = torch.sigmoid(next_esti_sdf * inv_s)
        alpha = (prev_cdf - next_cdf + 1e-5) / (prev_cdf + 1e-5)
        weights = alpha * torch.cumprod(
            torch.cat([torch.ones([batch_size, 1]), 1. - alpha + 1e-7], -1), -1)[:, :-1]

        z_samples = sample_pdf(z_vals, weights, n_importance, det=True).detach()
        return z_samples
//  [python]
def test_circle_boundary_grid_sampling_with_n():
    C = Circle(R2('x'), [0, 0], 4).boundary
    points = C.sample_grid(n=30)
    assert points.as_tensor.shape == (30, 2)
    assert all(torch.isclose(torch.norm(points.as_tensor, dim=1),
                             torch.tensor(4.0)))
//  [python]
def test_circle_boundary_grid_sampling_with_n():
    C = Circle(R2('x'), [0, 0], 4).boundary
    points = C.sample_grid(n=30)
    assert points.as_tensor.shape == (30, 2)
    assert all(torch.isclose(torch.linalg.norm(points.as_tensor, dim=1),
                             torch.tensor(4.0)))
//  [python]
def test_circle_boundary_grid_sampling_with_d():
    C = Circle(R2('x'), [0, 0], 4).boundary
    points = C.sample_grid(d=13)
    assert all(torch.isclose(torch.norm(points.as_tensor, dim=1),
                             torch.tensor(4.0)))
//  [python]
def test_circle_boundary_grid_sampling_with_d():
    C = Circle(R2('x'), [0, 0], 4).boundary
    points = C.sample_grid(d=13)
    assert all(torch.isclose(torch.linalg.norm(points.as_tensor, dim=1),
                             torch.tensor(4.0)))
//  [python]
def compute_jitter(x):
    """
    Compute jitter for the input tensor
    """
    return torch.norm(x[:, 2:] + x[:, :-2] - 2 * x[:, 1:-1], dim=-1)
//  [python]
def compute_jitter(x):
    """
    Compute jitter for the input tensor
    """
    return torch.linalg.norm(x[:, 2:] + x[:, :-2] - 2 * x[:, 1:-1], dim=-1)
//  [python]
def compute_sample_grads_sums(model, loss_fn, data, targets, additional_info, num_class, grad_epoch):
    """ manually process each sample with per sample gradient, naive implementation """
    if targets.max() >= num_class:
        raise ValueError("targets max value should be less than num_class")

    batch_size = data.shape[0]
    grad_sum = 0
    grad_sum_squared = 0
    grad_sum_clean = 0
    grad_sum_squared_clean = 0
    grad_sum_bd = 0
    grad_sum_squared_bd = 0
    grad_sum_class = [0 for _ in range(num_class)]
    grad_sum_squared_class = [0 for _ in range(num_class)]

    ori_idx, poi_indicator, ori_target = additional_info
    grad_dis_clean = []
    grad_dis_bd = []
    grad_dis_class = [[] for _ in range(num_class)]

    grad_cos_clean = []
    grad_cos_bd = []
    grad_cos_class = [[] for _ in range(num_class)]

    for i in range(batch_size):
        grad_i = compute_grad(model, loss_fn, data[i], targets[i])
        grad_sum += grad_i
        grad_sum_squared += grad_i.square()
        if poi_indicator[i] == 1:
            grad_sum_bd += grad_i
            grad_sum_squared_bd += grad_i.square()
            grad_dis_bd.append(torch.norm(grad_i.flatten()))
            grad_cos_bd.append(F.cosine_similarity(grad_i.flatten(), grad_epoch.flatten(),dim=0))
        else:
            grad_sum_clean += grad_i
            grad_sum_squared_clean += grad_i.square()
            grad_dis_clean.append(torch.norm(grad_i.flatten()))
            grad_cos_clean.append(F.cosine_similarity(grad_i.flatten(), grad_epoch.flatten(),dim=0))
            
            grad_sum_class[targets[i]] += grad_i
            grad_sum_squared_class[targets[i]] += grad_i.square()
            grad_dis_class[targets[i]].append(
                torch.norm(grad_i.flatten()))

            grad_cos_class[targets[i]].append(
                F.cosine_similarity(grad_i.flatten(), grad_epoch.flatten(),dim=0))

    # To avoid GPU memory overflow, we use numpy instead of torch.
    # grad_sum = grad_sum.cpu().numpy()
    # grad_sum_squared = grad_sum_squared.cpu().numpy()
    # grad_sum_clean = grad_sum_clean.cpu().numpy()
    # grad_sum_squared_clean = grad_sum_squared_clean.cpu().numpy()
    # grad_sum_bd = grad_sum_bd.cpu().numpy()
    # grad_sum_squared_bd = grad_sum_squared_bd.cpu().numpy()
    return grad_sum, grad_sum_squared, grad_sum_clean, grad_sum_squared_clean, grad_sum_bd, grad_sum_squared_bd, grad_dis_clean, grad_dis_bd, grad_sum_class, grad_sum_squared_class, grad_dis_class, grad_cos_clean, grad_cos_bd, grad_cos_class
//  [python]
def compute_sample_grads_sums(model, loss_fn, data, targets, additional_info, num_class, grad_epoch):
    """ manually process each sample with per sample gradient, naive implementation """
    if targets.max() >= num_class:
        raise ValueError("targets max value should be less than num_class")

    batch_size = data.shape[0]
    grad_sum = 0
    grad_sum_squared = 0
    grad_sum_clean = 0
    grad_sum_squared_clean = 0
    grad_sum_bd = 0
    grad_sum_squared_bd = 0
    grad_sum_class = [0 for _ in range(num_class)]
    grad_sum_squared_class = [0 for _ in range(num_class)]

    ori_idx, poi_indicator, ori_target = additional_info
    grad_dis_clean = []
    grad_dis_bd = []
    grad_dis_class = [[] for _ in range(num_class)]

    grad_cos_clean = []
    grad_cos_bd = []
    grad_cos_class = [[] for _ in range(num_class)]

    for i in range(batch_size):
        grad_i = compute_grad(model, loss_fn, data[i], targets[i])
        grad_sum += grad_i
        grad_sum_squared += grad_i.square()
        if poi_indicator[i] == 1:
            grad_sum_bd += grad_i
            grad_sum_squared_bd += grad_i.square()
            grad_dis_bd.append(torch.linalg.norm(grad_i.flatten()))
            grad_cos_bd.append(F.cosine_similarity(grad_i.flatten(), grad_epoch.flatten(),dim=0))
        else:
            grad_sum_clean += grad_i
            grad_sum_squared_clean += grad_i.square()
            grad_dis_clean.append(torch.linalg.norm(grad_i.flatten()))
            grad_cos_clean.append(F.cosine_similarity(grad_i.flatten(), grad_epoch.flatten(),dim=0))
            
            grad_sum_class[targets[i]] += grad_i
            grad_sum_squared_class[targets[i]] += grad_i.square()
            grad_dis_class[targets[i]].append(
                torch.linalg.norm(grad_i.flatten()))

            grad_cos_class[targets[i]].append(
                F.cosine_similarity(grad_i.flatten(), grad_epoch.flatten(),dim=0))

    # To avoid GPU memory overflow, we use numpy instead of torch.
    # grad_sum = grad_sum.cpu().numpy()
    # grad_sum_squared = grad_sum_squared.cpu().numpy()
    # grad_sum_clean = grad_sum_clean.cpu().numpy()
    # grad_sum_squared_clean = grad_sum_squared_clean.cpu().numpy()
    # grad_sum_bd = grad_sum_bd.cpu().numpy()
    # grad_sum_squared_bd = grad_sum_squared_bd.cpu().numpy()
    return grad_sum, grad_sum_squared, grad_sum_clean, grad_sum_squared_clean, grad_sum_bd, grad_sum_squared_bd, grad_dis_clean, grad_dis_bd, grad_sum_class, grad_sum_squared_class, grad_dis_class, grad_cos_clean, grad_cos_bd, grad_cos_class
//  [python]
def forward(self, x, y):
        """ x : B x T x C / T x C
            y : B x T x C / T x C
            return : B, B or 1, 1 [norm, error in log10] """
        if len(x.shape) > 2: # B x T x C
            padded_x = F.pad(x, (0, 0, 0, self.fftsize-x.shape[1]), "constant", 0)
            padded_y = F.pad(y, (0, 0, 0, self.fftsize-y.shape[1]), "constant", 0)

            csp_x = torch.fft.fftn(padded_x)
            csp_y = torch.fft.fftn(padded_y)

            magsp_x = torch.abs(csp_x)
            magsp_y = torch.abs(csp_y)

            diff = magsp_y - magsp_x
            norm = torch.norm(diff, 'fro', dim=(1,2)) / torch.norm(magsp_y, 'fro', dim=(1,2)) \
                    + diff.abs().sum(-1).sum(-1) / magsp_y.sum(-1).sum(-1)
            if x.shape[1] > 1:
                log_diff = torch.log10(torch.clamp(magsp_y, min=1e-13)) - torch.log10(torch.clamp(magsp_x, min=1e-13))
                err = log_diff.abs().mean(-1).mean(-1) + (log_diff**2).mean(-1).mean(-1).sqrt()
            else:
                err = (torch.log10(torch.clamp(magsp_y, min=1e-13)) - torch.log10(torch.clamp(magsp_x, min=1e-13))).abs().mean(-1).mean(-1)
        else: # T x C
            padded_x = F.pad(x, (0, self.fftsize-x.shape[1]), "constant", 0)
            padded_y = F.pad(y, (0, self.fftsize-y.shape[1]), "constant", 0)

            csp_x = torch.fft.fftn(padded_x)
            csp_y = torch.fft.fftn(padded_y)

            magsp_x = torch.abs(csp_x)
            magsp_y = torch.abs(csp_y)

            diff = magsp_y - magsp_x
            norm = torch.norm(diff, 'fro') / torch.norm(magsp_y, 'fro') + diff.abs().sum() / magsp_y.sum()
            if x.shape[0] > 1:
                log_diff = torch.log10(torch.clamp(magsp_y, min=1e-13)) - torch.log10(torch.clamp(magsp_x, min=1e-13))
                err = log_diff.abs().mean() + (log_diff**2).mean().sqrt()
            else:
                err = (torch.log10(torch.clamp(magsp_y, min=1e-13)) - torch.log10(torch.clamp(magsp_x, min=1e-13))).abs().mean()

        return norm, err
//  [python]
def forward(self, x, y):
        """ x : B x T x C / T x C
            y : B x T x C / T x C
            return : B, B or 1, 1 [norm, error in log10] """
        if len(x.shape) > 2: # B x T x C
            padded_x = F.pad(x, (0, 0, 0, self.fftsize-x.shape[1]), "constant", 0)
            padded_y = F.pad(y, (0, 0, 0, self.fftsize-y.shape[1]), "constant", 0)

            csp_x = torch.fft.fftn(padded_x)
            csp_y = torch.fft.fftn(padded_y)

            magsp_x = torch.abs(csp_x)
            magsp_y = torch.abs(csp_y)

            diff = magsp_y - magsp_x
            norm = LA.norm(diff, 'fro', dim=(1,2)) / LA.norm(magsp_y, 'fro', dim=(1,2)) \
                    + diff.abs().sum(-1).sum(-1) / magsp_y.sum(-1).sum(-1)
            if x.shape[1] > 1:
                log_diff = torch.log10(torch.clamp(magsp_y, min=1e-13)) - torch.log10(torch.clamp(magsp_x, min=1e-13))
                err = log_diff.abs().mean(-1).mean(-1) + (log_diff**2).mean(-1).mean(-1).sqrt()
            else:
                err = (torch.log10(torch.clamp(magsp_y, min=1e-13)) - torch.log10(torch.clamp(magsp_x, min=1e-13))).abs().mean(-1).mean(-1)
        else: # T x C
            padded_x = F.pad(x, (0, self.fftsize-x.shape[1]), "constant", 0)
            padded_y = F.pad(y, (0, self.fftsize-y.shape[1]), "constant", 0)

            csp_x = torch.fft.fftn(padded_x)
            csp_y = torch.fft.fftn(padded_y)

            magsp_x = torch.abs(csp_x)
            magsp_y = torch.abs(csp_y)

            diff = magsp_y - magsp_x
            norm = LA.norm(diff, 'fro') / LA.norm(magsp_y, 'fro') + diff.abs().sum() / magsp_y.sum()
            if x.shape[0] > 1:
                log_diff = torch.log10(torch.clamp(magsp_y, min=1e-13)) - torch.log10(torch.clamp(magsp_x, min=1e-13))
                err = log_diff.abs().mean() + (log_diff**2).mean().sqrt()
            else:
                err = (torch.log10(torch.clamp(magsp_y, min=1e-13)) - torch.log10(torch.clamp(magsp_x, min=1e-13))).abs().mean()

        return norm, err
//  [python]
def rel_diff(a, b=None, den=None):
    """
    Relative difference
    """
    if torch.is_tensor(a):
        order = 'fro'
        if den is None:
            d = min(torch.norm(a, p=order), torch.norm(b, p=order))
        else:
            b = den
            d = torch.norm(den, p=order)
        return torch.norm(a-den, p=order) / d
    else:
        if not type(a).__module__ == np.__name__:
            a = np.array([a])
            b = np.array([b])  
        if a.ndim == 0 or a.shape[0] == a.size or a.ndim>1 and a.shape[1] == a.size:
            order = None
        else:
            order = 'fro'
        if den is None:
            d = min(np.linalg.norm(a, ord=order), np.linalg.norm(b, ord=order))
        else:
            b = den
            d = np.linalg.norm(den, ord=order)
        return np.linalg.norm(a-b, ord=order) / d
//  [python]
def rel_diff(a, b=None, den=None):
    """
    Relative difference
    """
    if torch.is_tensor(a):
        order = 'fro'
        if den is None:
            d = min(torch.linalg.norm(a, ord=order), torch.linalg.norm(b, ord=order))
        else:
            b = den
            d = torch.linalg.norm(den, ord=order)
        return torch.linalg.norm(a-den, ord=order) / d
    else:
        if not type(a).__module__ == np.__name__:
            a = np.array([a])
            b = np.array([b])  
        if a.ndim == 0 or a.shape[0] == a.size or a.ndim>1 and a.shape[1] == a.size:
            order = None
        else:
            order = 'fro'
        if den is None:
            d = min(np.linalg.norm(a, ord=order), np.linalg.norm(b, ord=order))
        else:
            b = den
            d = np.linalg.norm(den, ord=order)
        return np.linalg.norm(a-b, ord=order) / d
//  [python]
def test_triangle_normals_if_corners_change():
    T = Triangle(R2('x'), origin, vec_1, [0, 1]).boundary
    points = torch.tensor([[0, 1], [0, 1], [2.0, 0.0], [2, 0.5], 
                           [0.5, 0.5], [0.0, 0.8]])
    points = Points(points, R2('x'))
    time = Points(torch.tensor([0, 1, 2.0, 3.0, 1.0, 0.0]).reshape(-1, 1), R1('t'))
    normals = T.normal(points, time)
    expected_normals = torch.tensor([[-0.3827,  0.9239], [-0.8112,  0.5847],
                                     [-0.2298, -0.9732], [ 0.2425,  0.9701],
                                     [-0.7071, -0.7071], [-1.0000,  0.0000]])
    assert normals.shape == (6, 2)
    one = torch.tensor(1.0)
    assert all(torch.isclose(torch.norm(normals, dim=1), one, atol=0.0001))
    assert torch.all(torch.isclose(expected_normals, normals, atol=0.0001))
//  [python]
def test_triangle_normals_if_corners_change():
    T = Triangle(R2('x'), origin, vec_1, [0, 1]).boundary
    points = torch.tensor([[0, 1], [0, 1], [2.0, 0.0], [2, 0.5], 
                           [0.5, 0.5], [0.0, 0.8]])
    points = Points(points, R2('x'))
    time = Points(torch.tensor([0, 1, 2.0, 3.0, 1.0, 0.0]).reshape(-1, 1), R1('t'))
    normals = T.normal(points, time)
    expected_normals = torch.tensor([[-0.3827,  0.9239], [-0.8112,  0.5847],
                                     [-0.2298, -0.9732], [ 0.2425,  0.9701],
                                     [-0.7071, -0.7071], [-1.0000,  0.0000]])
    assert normals.shape == (6, 2)
    one = torch.tensor(1.0)
    assert all(torch.isclose(torch.linalg.norm(normals, dim=1), one, atol=0.0001))
    assert torch.all(torch.isclose(expected_normals, normals, atol=0.0001))
//  [python]
def test_parallelogram_normals_if_corners_change():
    P = Parallelogram(R2('x'), origin, vec_1, [0, 1]).boundary
    points = torch.tensor([[0, 1], [0, 1], [2.0, 0.0],
                           [3.5, 1], [0.5, 0.5], [0.0, 0.8]])
    points = Points(points, R2('x'))
    time = Points(torch.tensor([0, 1, 2.0, 3.0, 1.0, 0.0]).reshape(-1, 1), R1('t'))
    normals = P.normal(points, time)
    expected_normals = torch.tensor([[-0.7071,  0.7071], [-0.9239,  0.3827],
                                     [-0.2298, -0.9732], [ 0.0000,  1.0000],
                                     [-0.7071, -0.7071], [-1.0000,  0.0000]])
    assert normals.shape == (6, 2)
    one = torch.tensor(1.0)
    assert all(torch.isclose(torch.norm(normals, dim=1), one, atol=0.0001))
    assert torch.all(torch.isclose(expected_normals, normals, atol=0.0001))
//  [python]
def test_parallelogram_normals_if_corners_change():
    P = Parallelogram(R2('x'), origin, vec_1, [0, 1]).boundary
    points = torch.tensor([[0, 1], [0, 1], [2.0, 0.0],
                           [3.5, 1], [0.5, 0.5], [0.0, 0.8]])
    points = Points(points, R2('x'))
    time = Points(torch.tensor([0, 1, 2.0, 3.0, 1.0, 0.0]).reshape(-1, 1), R1('t'))
    normals = P.normal(points, time)
    expected_normals = torch.tensor([[-0.7071,  0.7071], [-0.9239,  0.3827],
                                     [-0.2298, -0.9732], [ 0.0000,  1.0000],
                                     [-0.7071, -0.7071], [-1.0000,  0.0000]])
    assert normals.shape == (6, 2)
    one = torch.tensor(1.0)
    assert all(torch.isclose(torch.linalg.norm(normals, dim=1), one, atol=0.0001))
    assert torch.all(torch.isclose(expected_normals, normals, atol=0.0001))
//  [python]
def cal_bone_direction(parent_pose, children_pose):
    """ Calculate unit bone directions from adjacent 3D joints.
    
    Args:
    - parent_pose: (B, 1, 3)
    - children_pose: (B, Nc, 3), where Nc is the number of children.

    Returns:
    - orit: (B, 3, Nc), The 3D direction vector of bones.
    """

    orit = children_pose - parent_pose # broadcast
    orit = orit / (torch.norm(orit, dim=-1, keepdim=True) + 1e-8)
    orit = orit.permute(0, 2, 1)

    return orit
//  [python]
def cal_bone_direction(parent_pose, children_pose):
    """ Calculate unit bone directions from adjacent 3D joints.
    
    Args:
    - parent_pose: (B, 1, 3)
    - children_pose: (B, Nc, 3), where Nc is the number of children.

    Returns:
    - orit: (B, 3, Nc), The 3D direction vector of bones.
    """

    orit = children_pose - parent_pose # broadcast
    orit = orit / (torch.linalg.norm(orit, dim=-1, keepdim=True) + 1e-8)
    orit = orit.permute(0, 2, 1)

    return orit
//  [python]
def score(self, sample: torch.Tensor):
        # Evaluation of the score based on the energy
        with torch.no_grad():
            flat = sample - self.b_prime
            out = self.model(sample)
            energies = 0.5 * torch.sum(torch.square(flat), dim=1) - torch.sum(out, dim=1)

        # Evaluation of the score based on the reconstruction error
        sample.requires_grad_()
        out = self.model(sample)
        energy = self.energy(sample, out)
        dEn_dX = torch.autograd.grad(energy, sample)[0]
        rec_errs = torch.norm(dEn_dX, 2, keepdim=False, dim=1)
        return energies.cpu().numpy(), rec_errs.cpu().numpy()
//  [python]
def score(self, sample: torch.Tensor):
        # Evaluation of the score based on the energy
        with torch.no_grad():
            flat = sample - self.b_prime
            out = self.model(sample)
            energies = 0.5 * torch.sum(torch.square(flat), dim=1) - torch.sum(out, dim=1)

        # Evaluation of the score based on the reconstruction error
        sample.requires_grad_()
        out = self.model(sample)
        energy = self.energy(sample, out)
        dEn_dX = torch.autograd.grad(energy, sample)[0]
        rec_errs = torch.linalg.norm(dEn_dX, 2, keepdim=False, dim=1)
        return energies.cpu().numpy(), rec_errs.cpu().numpy()
//  [python]
def forward(self, output, params, input_keypoints, bbox, 
                reprojection_weight=100., regularize_weight=60.0, 
                consistency_weight=10.0, sprior_weight=0.04, 
                smooth_weight=20.0, sigma=100):
        
        pose, shape, cam = params
        scale = bbox[..., 2:].unsqueeze(-1) * 200.
        
        # Loss 1. Data term
        pred_keypoints = output.full_joints2d[..., :17, :]
        joints_conf = input_keypoints[..., -1:]
        reprojection_error = gmof(pred_keypoints - input_keypoints[..., :-1], sigma)
        reprojection_error = ((reprojection_error * joints_conf) / scale).mean()
        
        # Loss 2. Regularization term
        regularize_error = torch.norm(pose - self.init_pose, dim=-1).mean()
        
        # Loss 3. Shape prior and consistency error
        consistency_error = shape.std(dim=1).mean()
        sprior_error = torch.norm(shape, dim=-1).mean()
        shape_error = sprior_weight * sprior_error + consistency_weight * consistency_error
        
        # Loss 4. Smooth loss
        pose_diff = compute_jitter(pose).mean()
        cam_diff = compute_jitter(cam).mean()
        smooth_error = pose_diff + cam_diff
        
        # Sum up losses
        loss = {
            'reprojection': reprojection_weight * reprojection_error,
            'regularize': regularize_weight * regularize_error,
            'shape': shape_error,
            'smooth': smooth_weight * smooth_error
        }
        
        return loss
//  [python]
def forward(self, output, params, input_keypoints, bbox, 
                reprojection_weight=100., regularize_weight=60.0, 
                consistency_weight=10.0, sprior_weight=0.04, 
                smooth_weight=20.0, sigma=100):
        
        pose, shape, cam = params
        scale = bbox[..., 2:].unsqueeze(-1) * 200.
        
        # Loss 1. Data term
        pred_keypoints = output.full_joints2d[..., :17, :]
        joints_conf = input_keypoints[..., -1:]
        reprojection_error = gmof(pred_keypoints - input_keypoints[..., :-1], sigma)
        reprojection_error = ((reprojection_error * joints_conf) / scale).mean()
        
        # Loss 2. Regularization term
        regularize_error = torch.linalg.norm(pose - self.init_pose, dim=-1).mean()
        
        # Loss 3. Shape prior and consistency error
        consistency_error = shape.std(dim=1).mean()
        sprior_error = torch.linalg.norm(shape, dim=-1).mean()
        shape_error = sprior_weight * sprior_error + consistency_weight * consistency_error
        
        # Loss 4. Smooth loss
        pose_diff = compute_jitter(pose).mean()
        cam_diff = compute_jitter(cam).mean()
        smooth_error = pose_diff + cam_diff
        
        # Sum up losses
        loss = {
            'reprojection': reprojection_weight * reprojection_error,
            'regularize': regularize_weight * regularize_error,
            'shape': shape_error,
            'smooth': smooth_weight * smooth_error
        }
        
        return loss
//  [python]
def resolve_symmetry(xs, Rsnat_all, xsnat, Rsnat_all_alt, xsnat_alt, atm_mask):
    dists = torch.norm( xs[:,:,None,:] - xs[atm_mask,:][None,None,:,:], dim=-1)
    dists_nat = torch.norm( xsnat[:,:,None,:] - xsnat[atm_mask,:][None,None,:,:], dim=-1)
    dists_natalt = torch.norm( xsnat_alt[:,:,None,:] - xsnat_alt[atm_mask,:][None,None,:,:], dim=-1)

    drms_nat = torch.sum(torch.abs(dists_nat-dists),dim=(-1,-2))
    drms_natalt = torch.sum(torch.abs(dists_nat-dists_natalt), dim=(-1,-2))

    Rsnat_symm = Rsnat_all
    xs_symm = xsnat

    toflip = drms_natalt<drms_nat

    Rsnat_symm[toflip,...] = Rsnat_all_alt[toflip,...]
    xs_symm[toflip,...] = xsnat_alt[toflip,...]

    return Rsnat_symm, xs_symm
//  [python]
def resolve_symmetry(xs, Rsnat_all, xsnat, Rsnat_all_alt, xsnat_alt, atm_mask):
    dists = torch.linalg.norm( xs[:,:,None,:] - xs[atm_mask,:][None,None,:,:], dim=-1)
    dists_nat = torch.linalg.norm( xsnat[:,:,None,:] - xsnat[atm_mask,:][None,None,:,:], dim=-1)
    dists_natalt = torch.linalg.norm( xsnat_alt[:,:,None,:] - xsnat_alt[atm_mask,:][None,None,:,:], dim=-1)

    drms_nat = torch.sum(torch.abs(dists_nat-dists),dim=(-1,-2))
    drms_natalt = torch.sum(torch.abs(dists_nat-dists_natalt), dim=(-1,-2))

    Rsnat_symm = Rsnat_all
    xs_symm = xsnat

    toflip = drms_natalt<drms_nat

    Rsnat_symm[toflip,...] = Rsnat_all_alt[toflip,...]
    xs_symm[toflip,...] = xsnat_alt[toflip,...]

    return Rsnat_symm, xs_symm
//  [python]
def cal_skel_direction(joints, children, parents, num_joint, mt_child_joints, 
                       exclude_root=False, leaf_idx=None):
    """ Calculate unit bone directions of the whole skeleton from 3D joints.
    
    Args:    
    - joints: (B, J, 3), 3D joints.
    - mt_child_joints: list, index of joints with multiple children.
    - exclude_root: bool, if True then the root orientation is not calculated.
    - leaf_idx: list, index of leaf joints.
    
    Returns:
    - st_orit_list: (B, N_single 3, 1), where N_single is the number of single-child node.
    - mt_orit_list: list, each element is (B, 3, Nc), where Nc is the number of its children.
    """

    eps = 1e-8

    sg_child_joints = []
    for j in range(num_joint):
        if j in mt_child_joints: continue
        if leaf_idx is not None and j in leaf_idx: continue
        sg_child_joints.append(j)

    # single-child node
    sg_orit_list = [] 
    for pa in sg_child_joints:
        ch = children[pa]
        sg_pa_orit = joints[:, ch, :] - joints[:, pa, :]
        unit_sg_pa_orit = sg_pa_orit / (torch.norm(sg_pa_orit, dim=-1, keepdim=True) + eps)
        sg_orit_list.append(unit_sg_pa_orit)
    sg_orit_list = torch.stack(sg_orit_list, dim=0).transpose(0, 1).unsqueeze(-1)
    
    # multi-child node
    if exclude_root:
        new_mt_child_joints = mt_child_joints[1:]
    else:
        new_mt_child_joints = mt_child_joints
    mt_orit_list = [[] for i in range(len(new_mt_child_joints))]
    for i, pa in enumerate(new_mt_child_joints):
        for ch, map_pa in enumerate(parents):
            if pa == map_pa:
                mt_pa_orit = joints[:, ch, :] - joints[:, pa, :]
                unit_mt_pa_orit = mt_pa_orit / (torch.norm(mt_pa_orit, dim=-1, keepdim=True) + eps)
                mt_orit_list[i].append(unit_mt_pa_orit)

    mt_orit_list = [torch.stack(imt, dim=0).permute(1, 2, 0) for imt in mt_orit_list]  # list of (B, 3, Nc)
    
    return sg_orit_list, mt_orit_list
//  [python]
def cal_skel_direction(joints, children, parents, num_joint, mt_child_joints, 
                       exclude_root=False, leaf_idx=None):
    """ Calculate unit bone directions of the whole skeleton from 3D joints.
    
    Args:    
    - joints: (B, J, 3), 3D joints.
    - mt_child_joints: list, index of joints with multiple children.
    - exclude_root: bool, if True then the root orientation is not calculated.
    - leaf_idx: list, index of leaf joints.
    
    Returns:
    - st_orit_list: (B, N_single 3, 1), where N_single is the number of single-child node.
    - mt_orit_list: list, each element is (B, 3, Nc), where Nc is the number of its children.
    """

    eps = 1e-8

    sg_child_joints = []
    for j in range(num_joint):
        if j in mt_child_joints: continue
        if leaf_idx is not None and j in leaf_idx: continue
        sg_child_joints.append(j)

    # single-child node
    sg_orit_list = [] 
    for pa in sg_child_joints:
        ch = children[pa]
        sg_pa_orit = joints[:, ch, :] - joints[:, pa, :]
        unit_sg_pa_orit = sg_pa_orit / (torch.linalg.norm(sg_pa_orit, dim=-1, keepdim=True) + eps)
        sg_orit_list.append(unit_sg_pa_orit)
    sg_orit_list = torch.stack(sg_orit_list, dim=0).transpose(0, 1).unsqueeze(-1)
    
    # multi-child node
    if exclude_root:
        new_mt_child_joints = mt_child_joints[1:]
    else:
        new_mt_child_joints = mt_child_joints
    mt_orit_list = [[] for i in range(len(new_mt_child_joints))]
    for i, pa in enumerate(new_mt_child_joints):
        for ch, map_pa in enumerate(parents):
            if pa == map_pa:
                mt_pa_orit = joints[:, ch, :] - joints[:, pa, :]
                unit_mt_pa_orit = mt_pa_orit / (torch.linalg.norm(mt_pa_orit, dim=-1, keepdim=True) + eps)
                mt_orit_list[i].append(unit_mt_pa_orit)

    mt_orit_list = [torch.stack(imt, dim=0).permute(1, 2, 0) for imt in mt_orit_list]  # list of (B, 3, Nc)
    
    return sg_orit_list, mt_orit_list
//  [python]
def forward(
        self,
        image_embeds: torch.FloatTensor,
        query_embeds: torch.FloatTensor,
        query_mask: torch.Tensor,
    ) -> Tuple[torch.FloatTensor]:

        image_class_embeds = self.dense0(image_embeds)

        # Normalize image and text features
        image_class_embeds /= torch.norm(image_class_embeds, dim=-1, keepdim=True) + 1e-6
        query_embeds /= torch.norm(query_embeds, dim=-1, keepdim=True) + 1e-6

        # Get class predictions
        pred_logits = torch.einsum("...pd,...qd->...pq", image_class_embeds, query_embeds)

        # Apply a learnable shift and scale to logits
        logit_shift = self.logit_shift(image_embeds)
        logit_scale = self.logit_scale(image_embeds)
        logit_scale = self.elu(logit_scale) + 1
        pred_logits = (pred_logits + logit_shift) * logit_scale

        if query_mask is not None:
            if query_mask.ndim > 1:
                query_mask = torch.unsqueeze(query_mask, dim=-2)

            pred_logits = pred_logits.to(torch.float64)
            pred_logits = torch.where(query_mask == 0, -1e6, pred_logits)
            pred_logits = pred_logits.to(torch.float32)

        return (pred_logits, image_class_embeds)
//  [python]
def forward(
        self,
        image_embeds: torch.FloatTensor,
        query_embeds: torch.FloatTensor,
        query_mask: torch.Tensor,
    ) -> Tuple[torch.FloatTensor]:

        image_class_embeds = self.dense0(image_embeds)

        # Normalize image and text features
        image_class_embeds /= torch.linalg.norm(image_class_embeds, dim=-1, keepdim=True) + 1e-6
        query_embeds /= torch.linalg.norm(query_embeds, dim=-1, keepdim=True) + 1e-6

        # Get class predictions
        pred_logits = torch.einsum("...pd,...qd->...pq", image_class_embeds, query_embeds)

        # Apply a learnable shift and scale to logits
        logit_shift = self.logit_shift(image_embeds)
        logit_scale = self.logit_scale(image_embeds)
        logit_scale = self.elu(logit_scale) + 1
        pred_logits = (pred_logits + logit_shift) * logit_scale

        if query_mask is not None:
            if query_mask.ndim > 1:
                query_mask = torch.unsqueeze(query_mask, dim=-2)

            pred_logits = pred_logits.to(torch.float64)
            pred_logits = torch.where(query_mask == 0, -1e6, pred_logits)
            pred_logits = pred_logits.to(torch.float32)

        return (pred_logits, image_class_embeds)
//  [python]
def get_geometric_structure_embeddings(points, angle_k=3):

    batch_size, num_point, _ = points.shape

    dist_map = torch.sqrt(square_distance(points, points))  # (B, N, N)

    knn_indices = dist_map.topk(k=angle_k + 1, dim=2, largest=False)[1]  # (B, N, k)
    knn_indices = knn_indices[:, :, 1:]
    knn_indices = knn_indices.unsqueeze(3).expand(batch_size, num_point, angle_k, 3)  # (B, N, k, 3)
    expanded_points = points.unsqueeze(1).expand(batch_size, num_point, num_point, 3)  # (B, N, N, 3)
    knn_points = torch.gather(expanded_points, dim=2, index=knn_indices)  # (B, N, k, 3)
    ref_vectors = knn_points - points.unsqueeze(2)  # (B, N, k, 3)
    anc_vectors = points.unsqueeze(1) - points.unsqueeze(2)  # (B, N, N, 3)
    ref_vectors = ref_vectors.unsqueeze(2).expand(batch_size, num_point, num_point, angle_k, 3)  # (B, N, N, k, 3)
    anc_vectors = anc_vectors.unsqueeze(3).expand(batch_size, num_point, num_point, angle_k, 3)  # (B, N, N, k, 3)
    sin_values = torch.norm(torch.cross(ref_vectors, anc_vectors, dim=-1), dim=-1)  # (B, N, N, k)
    cos_values = torch.sum(ref_vectors * anc_vectors, dim=-1)  # (B, N, N, k)
    angles = torch.atan2(sin_values, cos_values)  # (B, N, N, k)

    return dist_map, angles
//  [python]
def get_geometric_structure_embeddings(points, angle_k=3):

    batch_size, num_point, _ = points.shape

    dist_map = torch.sqrt(square_distance(points, points))  # (B, N, N)

    knn_indices = dist_map.topk(k=angle_k + 1, dim=2, largest=False)[1]  # (B, N, k)
    knn_indices = knn_indices[:, :, 1:]
    knn_indices = knn_indices.unsqueeze(3).expand(batch_size, num_point, angle_k, 3)  # (B, N, k, 3)
    expanded_points = points.unsqueeze(1).expand(batch_size, num_point, num_point, 3)  # (B, N, N, 3)
    knn_points = torch.gather(expanded_points, dim=2, index=knn_indices)  # (B, N, k, 3)
    ref_vectors = knn_points - points.unsqueeze(2)  # (B, N, k, 3)
    anc_vectors = points.unsqueeze(1) - points.unsqueeze(2)  # (B, N, N, 3)
    ref_vectors = ref_vectors.unsqueeze(2).expand(batch_size, num_point, num_point, angle_k, 3)  # (B, N, N, k, 3)
    anc_vectors = anc_vectors.unsqueeze(3).expand(batch_size, num_point, num_point, angle_k, 3)  # (B, N, N, k, 3)
    sin_values = torch.linalg.norm(torch.cross(ref_vectors, anc_vectors, dim=-1), dim=-1)  # (B, N, N, k)
    cos_values = torch.sum(ref_vectors * anc_vectors, dim=-1)  # (B, N, N, k)
    angles = torch.atan2(sin_values, cos_values)  # (B, N, N, k)

    return dist_map, angles
//  [python]
def _encode_prompt(
        self,
        prompt,
        device,
        num_images_per_prompt,
        do_classifier_free_guidance,
    ):
        len(prompt) if isinstance(prompt, list) else 1

        # YiYi Notes: set pad_token_id to be 0, not sure why I can't set in the config file
        self.tokenizer.pad_token_id = 0
        # get prompt text embeddings
        text_inputs = self.tokenizer(
            prompt,
            padding="max_length",
            max_length=self.tokenizer.model_max_length,
            truncation=True,
            return_tensors="pt",
        )
        text_input_ids = text_inputs.input_ids
        untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids

        if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
            removed_text = self.tokenizer.batch_decode(untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1])
            logger.warning(
                "The following part of your input was truncated because CLIP can only handle sequences up to"
                f" {self.tokenizer.model_max_length} tokens: {removed_text}"
            )

        text_encoder_output = self.text_encoder(text_input_ids.to(device))
        prompt_embeds = text_encoder_output.text_embeds

        prompt_embeds = prompt_embeds.repeat_interleave(num_images_per_prompt, dim=0)
        # in Shap-E it normalize the prompt_embeds and then later rescale it
        prompt_embeds = prompt_embeds / torch.norm(prompt_embeds, dim=-1, keepdim=True)

        if do_classifier_free_guidance:
            negative_prompt_embeds = torch.zeros_like(prompt_embeds)

            # For classifier free guidance, we need to do two forward passes.
            # Here we concatenate the unconditional and text embeddings into a single batch
            # to avoid doing two forward passes
            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])

        # Rescale the features to have unit variance
        prompt_embeds = math.sqrt(prompt_embeds.shape[1]) * prompt_embeds

        return prompt_embeds
//  [python]
def _encode_prompt(
        self,
        prompt,
        device,
        num_images_per_prompt,
        do_classifier_free_guidance,
    ):
        len(prompt) if isinstance(prompt, list) else 1

        # YiYi Notes: set pad_token_id to be 0, not sure why I can't set in the config file
        self.tokenizer.pad_token_id = 0
        # get prompt text embeddings
        text_inputs = self.tokenizer(
            prompt,
            padding="max_length",
            max_length=self.tokenizer.model_max_length,
            truncation=True,
            return_tensors="pt",
        )
        text_input_ids = text_inputs.input_ids
        untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids

        if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
            removed_text = self.tokenizer.batch_decode(untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1])
            logger.warning(
                "The following part of your input was truncated because CLIP can only handle sequences up to"
                f" {self.tokenizer.model_max_length} tokens: {removed_text}"
            )

        text_encoder_output = self.text_encoder(text_input_ids.to(device))
        prompt_embeds = text_encoder_output.text_embeds

        prompt_embeds = prompt_embeds.repeat_interleave(num_images_per_prompt, dim=0)
        # in Shap-E it normalize the prompt_embeds and then later rescale it
        prompt_embeds = prompt_embeds / torch.linalg.norm(prompt_embeds, dim=-1, keepdim=True)

        if do_classifier_free_guidance:
            negative_prompt_embeds = torch.zeros_like(prompt_embeds)

            # For classifier free guidance, we need to do two forward passes.
            # Here we concatenate the unconditional and text embeddings into a single batch
            # to avoid doing two forward passes
            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])

        # Rescale the features to have unit variance
        prompt_embeds = math.sqrt(prompt_embeds.shape[1]) * prompt_embeds

        return prompt_embeds
//  [python]
def compute_action(self, observation: torch.Tensor, u_range: float) -> torch.Tensor:
        assert self.continuous_actions

        # First calculate the closest point to a circle of radius circle_radius given the current position
        circle_origin = torch.zeros(1, 2, device=observation.device)
        circle_radius = 0.75
        current_pos = observation[:, :2]
        v = current_pos - circle_origin
        closest_point_on_circ = (
            circle_origin + v / torch.norm(v, dim=1).unsqueeze(1) * circle_radius
        )

        # calculate the normal vector of the vector from the origin of the circle to that closest point
        # on the circle. Adding this scaled normal vector to the other vector gives us a target point we
        # try to reach, thus resulting in a circular motion.
        closest_point_on_circ_normal = torch.stack(
            [closest_point_on_circ[:, Y], -closest_point_on_circ[:, X]], dim=1
        )
        closest_point_on_circ_normal /= torch.norm(
            closest_point_on_circ_normal, dim=1
        ).unsqueeze(1)
        closest_point_on_circ_normal *= 0.1
        des_pos = closest_point_on_circ + closest_point_on_circ_normal

        # Move away from other agents within visibility range
        lidar_agents = observation[:, 4:16]
        agent_visible = torch.any(lidar_agents < 0.15, dim=1)
        _, agent_dir_index = torch.min(lidar_agents, dim=1)
        agent_dir = agent_dir_index / lidar_agents.shape[1] * 2 * torch.pi
        agent_vec = torch.stack([torch.cos(agent_dir), torch.sin(agent_dir)], dim=1)
        des_pos_agent = current_pos - agent_vec * 0.1
        des_pos[agent_visible] = des_pos_agent[agent_visible]

        # Move towards targets within visibility range
        lidar_targets = observation[:, 16:28]
        target_visible = torch.any(lidar_targets < 0.3, dim=1)
        _, target_dir_index = torch.min(lidar_targets, dim=1)
        target_dir = target_dir_index / lidar_targets.shape[1] * 2 * torch.pi
        target_vec = torch.stack([torch.cos(target_dir), torch.sin(target_dir)], dim=1)
        des_pos_target = current_pos + target_vec * 0.1
        des_pos[target_visible] = des_pos_target[target_visible]

        action = torch.clamp(
            (des_pos - current_pos) * 10,
            min=-u_range,
            max=u_range,
        )

        return action
//  [python]
def compute_action(self, observation: torch.Tensor, u_range: float) -> torch.Tensor:
        assert self.continuous_actions

        # First calculate the closest point to a circle of radius circle_radius given the current position
        circle_origin = torch.zeros(1, 2, device=observation.device)
        circle_radius = 0.75
        current_pos = observation[:, :2]
        v = current_pos - circle_origin
        closest_point_on_circ = (
            circle_origin + v / torch.linalg.norm(v, dim=1).unsqueeze(1) * circle_radius
        )

        # calculate the normal vector of the vector from the origin of the circle to that closest point
        # on the circle. Adding this scaled normal vector to the other vector gives us a target point we
        # try to reach, thus resulting in a circular motion.
        closest_point_on_circ_normal = torch.stack(
            [closest_point_on_circ[:, Y], -closest_point_on_circ[:, X]], dim=1
        )
        closest_point_on_circ_normal /= torch.linalg.norm(
            closest_point_on_circ_normal, dim=1
        ).unsqueeze(1)
        closest_point_on_circ_normal *= 0.1
        des_pos = closest_point_on_circ + closest_point_on_circ_normal

        # Move away from other agents within visibility range
        lidar_agents = observation[:, 4:16]
        agent_visible = torch.any(lidar_agents < 0.15, dim=1)
        _, agent_dir_index = torch.min(lidar_agents, dim=1)
        agent_dir = agent_dir_index / lidar_agents.shape[1] * 2 * torch.pi
        agent_vec = torch.stack([torch.cos(agent_dir), torch.sin(agent_dir)], dim=1)
        des_pos_agent = current_pos - agent_vec * 0.1
        des_pos[agent_visible] = des_pos_agent[agent_visible]

        # Move towards targets within visibility range
        lidar_targets = observation[:, 16:28]
        target_visible = torch.any(lidar_targets < 0.3, dim=1)
        _, target_dir_index = torch.min(lidar_targets, dim=1)
        target_dir = target_dir_index / lidar_targets.shape[1] * 2 * torch.pi
        target_vec = torch.stack([torch.cos(target_dir), torch.sin(target_dir)], dim=1)
        des_pos_target = current_pos + target_vec * 0.1
        des_pos[target_visible] = des_pos_target[target_visible]

        action = torch.clamp(
            (des_pos - current_pos) * 10,
            min=-u_range,
            max=u_range,
        )

        return action
//  [python]
def SE3_logmap(T, eps=1e-6):
    w = SO3_logmap(T[:, :3, :3])
    theta = torch.norm(w, dim=(1))
    theta = torch.clamp(theta, min=eps)
    w_norm = w / theta
    tan = torch.tan(0.5 * theta)

    t = T[:, :3, 3]
    wnorm_x_t = torch.cross(w_norm, t)
    V_inv_t = (
        t
        - (0.5 * t) * wnorm_x_t
        + (1.0 - theta / (2.0 * tan)) * torch.cross(w_norm, wnorm_x_t)
    )
    xi = torch.cat((w, V_inv_t), dim=-1)

    return xi
//  [python]
def SE3_logmap(T, eps=1e-6):
    w = SO3_logmap(T[:, :3, :3])
    theta = torch.linalg.norm(w, dim=(1))
    theta = torch.clamp(theta, min=eps)
    w_norm = w / theta
    tan = torch.tan(0.5 * theta)

    t = T[:, :3, 3]
    wnorm_x_t = torch.cross(w_norm, t)
    V_inv_t = (
        t
        - (0.5 * t) * wnorm_x_t
        + (1.0 - theta / (2.0 * tan)) * torch.cross(w_norm, wnorm_x_t)
    )
    xi = torch.cat((w, V_inv_t), dim=-1)

    return xi