From 561e121a42461f24840226400009594509b6b3fc Mon Sep 17 00:00:00 2001
From: brook-bee <1215918085@qq.com>
Date: Fri, 12 Jan 2024 15:48:36 +0800
Subject: [PATCH] add isaacgym_utils

---
 legged_gym/utils/isaacgym_utils.py | 584 +++++++++++++++++++++++++++++
 1 file changed, 584 insertions(+)
 create mode 100644 legged_gym/utils/isaacgym_utils.py

diff --git a/legged_gym/utils/isaacgym_utils.py b/legged_gym/utils/isaacgym_utils.py
new file mode 100644
index 0000000..3fd8c0c
--- /dev/null
+++ b/legged_gym/utils/isaacgym_utils.py
@@ -0,0 +1,584 @@
+import os
+import numpy as np
+import random
+import torch
+
+
+def set_seed(seed):
+    if seed == -1:
+        seed = np.random.randint(0, 10000)
+    print("Setting seed: {}".format(seed))
+
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    os.environ['PYTHONHASHSEED'] = str(seed)
+    torch.cuda.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+
+def to_torch(x, dtype=torch.float, device='cuda:0', requires_grad=False):
+    return torch.tensor(x,
+                        dtype=dtype,
+                        device=device,
+                        requires_grad=requires_grad)
+
+
+@torch.jit.script
+def quat_mul(a, b):
+    assert a.shape == b.shape
+    shape = a.shape
+    a = a.reshape(-1, 4)
+    b = b.reshape(-1, 4)
+
+    x1, y1, z1, w1 = a[:, 0], a[:, 1], a[:, 2], a[:, 3]
+    x2, y2, z2, w2 = b[:, 0], b[:, 1], b[:, 2], b[:, 3]
+    ww = (z1 + x1) * (x2 + y2)
+    yy = (w1 - y1) * (w2 + z2)
+    zz = (w1 + y1) * (w2 - z2)
+    xx = ww + yy + zz
+    qq = 0.5 * (xx + (z1 - x1) * (x2 - y2))
+    w = qq - ww + (z1 - y1) * (y2 - z2)
+    x = qq - xx + (x1 + w1) * (x2 + w2)
+    y = qq - yy + (w1 - x1) * (y2 + z2)
+    z = qq - zz + (z1 + y1) * (w2 - x2)
+
+    quat = torch.stack([x, y, z, w], dim=-1).view(shape)
+
+    return quat
+
+
+@torch.jit.script
+def normalize(x, eps: float = 1e-9):
+    return x / x.norm(p=2, dim=-1).clamp(min=eps, max=None).unsqueeze(-1)
+
+
+@torch.jit.script
+def quat_apply(a, b):
+    shape = b.shape
+    a = a.reshape(-1, 4)
+    b = b.reshape(-1, 3)
+    xyz = a[:, :3]
+    t = xyz.cross(b, dim=-1) * 2
+    return (b + a[:, 3:] * t + xyz.cross(t, dim=-1)).view(shape)
+
+
+@torch.jit.script
+def quat_rotate(q, v):
+    shape = q.shape
+    q_w = q[:, -1]
+    q_vec = q[:, :3]
+    a = v * (2.0 * q_w ** 2 - 1.0).unsqueeze(-1)
+    b = torch.cross(q_vec, v, dim=-1) * q_w.unsqueeze(-1) * 2.0
+    c = q_vec * \
+        torch.bmm(q_vec.view(shape[0], 1, 3), v.view(
+            shape[0], 3, 1)).squeeze(-1) * 2.0
+    return a + b + c
+
+
+@torch.jit.script
+def quat_rotate_inverse(q, v):
+    shape = q.shape
+    q_w = q[:, -1]
+    q_vec = q[:, :3]
+    a = v * (2.0 * q_w ** 2 - 1.0).unsqueeze(-1)
+    b = torch.cross(q_vec, v, dim=-1) * q_w.unsqueeze(-1) * 2.0
+    c = q_vec * \
+        torch.bmm(q_vec.view(shape[0], 1, 3), v.view(
+            shape[0], 3, 1)).squeeze(-1) * 2.0
+    return a - b + c
+
+
+@torch.jit.script
+def quat_conjugate(a):
+    shape = a.shape
+    a = a.reshape(-1, 4)
+    return torch.cat((-a[:, :3], a[:, -1:]), dim=-1).view(shape)
+
+
+@torch.jit.script
+def quat_unit(a):
+    return normalize(a)
+
+
+@torch.jit.script
+def quat_from_angle_axis(angle, axis):
+    theta = (angle / 2).unsqueeze(-1)
+    xyz = normalize(axis) * theta.sin()
+    w = theta.cos()
+    return quat_unit(torch.cat([xyz, w], dim=-1))
+
+
+@torch.jit.script
+def normalize_angle(x):
+    return torch.atan2(torch.sin(x), torch.cos(x))
+
+
+@torch.jit.script
+def tf_inverse(q, t):
+    q_inv = quat_conjugate(q)
+    return q_inv, -quat_apply(q_inv, t)
+
+
+@torch.jit.script
+def tf_apply(q, t, v):
+    return quat_apply(q, v) + t
+
+
+@torch.jit.script
+def tf_vector(q, v):
+    return quat_apply(q, v)
+
+
+@torch.jit.script
+def tf_combine(q1, t1, q2, t2):
+    return quat_mul(q1, q2), quat_apply(q1, t2) + t1
+
+
+@torch.jit.script
+def get_basis_vector(q, v):
+    return quat_rotate(q, v)
+
+
+def get_axis_params(value, axis_idx, x_value=0., dtype=np.float32, n_dims=3):
+    """construct arguments to `Vec` according to axis index.
+    """
+    zs = np.zeros((n_dims,))
+    assert axis_idx < n_dims, "the axis dim should be within the vector dimensions"
+    zs[axis_idx] = 1.
+    params = np.where(zs == 1., value, zs)
+    params[0] = x_value
+    return list(params.astype(dtype))
+
+
+@torch.jit.script
+def copysign(a, b):
+    # type: (float, Tensor) -> Tensor
+    a = torch.tensor(a, device=b.device, dtype=torch.float).repeat(b.shape[0])
+    return torch.abs(a) * torch.sign(b)
+
+
+@torch.jit.script
+def get_euler_xyz(q):
+    qx, qy, qz, qw = 0, 1, 2, 3
+    # roll (x-axis rotation)
+    sinr_cosp = 2.0 * (q[:, qw] * q[:, qx] + q[:, qy] * q[:, qz])
+    cosr_cosp = q[:, qw] * q[:, qw] - q[:, qx] * \
+                q[:, qx] - q[:, qy] * q[:, qy] + q[:, qz] * q[:, qz]
+    roll = torch.atan2(sinr_cosp, cosr_cosp)
+
+    # pitch (y-axis rotation)
+    sinp = 2.0 * (q[:, qw] * q[:, qy] - q[:, qz] * q[:, qx])
+    pitch = torch.where(
+        torch.abs(sinp) >= 1, copysign(np.pi / 2.0, sinp), torch.asin(sinp))
+
+    # yaw (z-axis rotation)
+    siny_cosp = 2.0 * (q[:, qw] * q[:, qz] + q[:, qx] * q[:, qy])
+    cosy_cosp = q[:, qw] * q[:, qw] + q[:, qx] * \
+                q[:, qx] - q[:, qy] * q[:, qy] - q[:, qz] * q[:, qz]
+    yaw = torch.atan2(siny_cosp, cosy_cosp)
+
+    return roll % (2 * np.pi), pitch % (2 * np.pi), yaw % (2 * np.pi)
+
+
+@torch.jit.script
+def quat_from_euler_xyz(roll, pitch, yaw):
+    cy = torch.cos(yaw * 0.5)
+    sy = torch.sin(yaw * 0.5)
+    cr = torch.cos(roll * 0.5)
+    sr = torch.sin(roll * 0.5)
+    cp = torch.cos(pitch * 0.5)
+    sp = torch.sin(pitch * 0.5)
+
+    qw = cy * cr * cp + sy * sr * sp
+    qx = cy * sr * cp - sy * cr * sp
+    qy = cy * cr * sp + sy * sr * cp
+    qz = sy * cr * cp - cy * sr * sp
+
+    return torch.stack([qx, qy, qz, qw], dim=-1)
+
+
+@torch.jit.script
+def torch_rand_float(lower, upper, shape, device):
+    # type: (float, float, Tuple[int, int], str) -> Tensor
+    return (upper - lower) * torch.rand(*shape, device=device) + lower
+
+
+@torch.jit.script
+def torch_random_dir_2(shape, device):
+    # type: (Tuple[int, int], str) -> Tensor
+    angle = torch_rand_float(-np.pi, np.pi, shape, device).squeeze(-1)
+    return torch.stack([torch.cos(angle), torch.sin(angle)], dim=-1)
+
+
+@torch.jit.script
+def tensor_clamp(t, min_t, max_t):
+    return torch.max(torch.min(t, max_t), min_t)
+
+
+@torch.jit.script
+def scale(x, lower, upper):
+    return (0.5 * (x + 1.0) * (upper - lower) + lower)
+
+
+@torch.jit.script
+def unscale(x, lower, upper):
+    return (2.0 * x - upper - lower) / (upper - lower)
+
+
+def unscale_np(x, lower, upper):
+    return (2.0 * x - upper - lower) / (upper - lower)
+
+
+def slope_platform_stairs_terrain(terrain,
+                                  slope=1,
+                                  step_width=0.2,
+                                  step_height=0.1,
+                                  num_steps=5):
+    """
+    Generate a sloped followed by a platform and a downstair
+
+    Parameters:
+        terrain (SubTerrain): the terrain
+        slope (int): positive or negative slope
+        step_width (float): the width of the step [meters]
+        step_height (float): the height of the step [meters]
+        num_steps (int): number of the step of the stair
+    Returns:
+        terrain (SubTerrain): update terrain
+    """
+    # switch parameters to discrete units
+    step_width = round(step_width / terrain.horizontal_scale)
+    step_height = round(step_height / terrain.vertical_scale)
+
+    height = step_height
+    for i in range(num_steps):
+        terrain.height_field_raw[i * step_width:(i + 1) *
+                                                step_width, :] += height
+        height += step_height
+
+    platform_height = height
+    slope_width = int(platform_height /
+                      (np.abs(slope) *
+                       (terrain.horizontal_scale / terrain.vertical_scale)))
+    x = np.arange(0, slope_width)
+    y = np.arange(0, terrain.length)
+    xx, yy = np.meshgrid(x, y, sparse=True)
+    xx = xx.reshape(slope_width, 1)
+    platform_width = int(terrain.width - num_steps * step_width - slope_width)
+    terrain.height_field_raw[num_steps * step_width:,
+    np.arange(terrain.length)] = platform_height
+    terrain.height_field_raw[-slope_width:,
+    np.arange(terrain.length)] -= (
+            platform_height * xx / slope_width).astype(
+        terrain.height_field_raw.dtype)
+    terrain.height_field_raw = terrain.height_field_raw[::-1]
+
+    return terrain
+
+
+def stairs_platform_slope_terrain(terrain,
+                                  step_width=0.2,
+                                  step_height=0.1,
+                                  num_steps=10,
+                                  slope=0.26):
+    """
+    Generate a stairs followed with a flat plane and a downslope
+
+    Parameters:
+        terrain (terrain): the terrain
+        step_width (float): the width of the step [meters]
+        step_height (float): the height of the step [meters]
+        num_steps (int): number of the step of the stair
+        slope (float): positive or negative slope
+    Returns:
+        terrain (SubTerrain): update terrain
+    """
+    # switch parameters to discrete units
+    step_width = round(step_width / terrain.horizontal_scale)
+    step_height = round(step_height / terrain.vertical_scale)
+
+    height = step_height
+    for i in range(num_steps):
+        terrain.height_field_raw[i * step_width:(i + 1) *
+                                                step_width, :] += height
+        height += step_height
+
+    platform_height = height
+    slope_width = int(platform_height /
+                      (np.abs(slope) *
+                       (terrain.horizontal_scale / terrain.vertical_scale)))
+    x = np.arange(0, slope_width)
+    y = np.arange(0, terrain.length)
+    xx, yy = np.meshgrid(x, y, sparse=True)
+    xx = xx.reshape(slope_width, 1)
+    platform_width = int(terrain.width - num_steps * step_width - slope_width)
+    terrain.height_field_raw[num_steps * step_width:,
+    np.arange(terrain.length)] = platform_height
+    terrain.height_field_raw[-slope_width:,
+    np.arange(terrain.length)] -= (
+            platform_height * xx / slope_width).astype(
+        terrain.height_field_raw.dtype)
+
+    return terrain
+
+
+if torch.cuda.is_available():
+    import string
+    import cupy as cp
+
+
+    def compute_meshes_normals_kernel():
+        compute_meshes_normals_kernel = cp.ElementwiseKernel(
+            in_params="raw U vertices, raw T triangles",
+            out_params="raw U output",
+            preamble=string.Template("""
+                __device__ bool cross(float p1x, float p1y, float p1z,
+                                        float p2x, float p2y, float p2z,
+                                        float& normalx,float& normaly,float& normalz)
+                {
+                    float nx = p1y*p2z - p1z*p2y;
+                    float ny = p1z*p2x - p1x*p2z;
+                    float nz = p1x*p2y - p1y*p2x;
+                    float l = sqrt(nx*nx + ny*ny + nz*nz);
+                    if (l <1e-4) {
+                        l = 1;
+                    }
+                    if (nz < 0.) {
+                        nx *=-1;
+                        ny *=-1;
+                        nz *=-1;
+                    }
+                    atomicAdd(&normalx, nx/l);
+                    atomicAdd(&normaly, ny/l);
+                    atomicAdd(&normalz, nz/l);
+
+                    return true;
+                }
+                """).substitute(),
+            operation=string.Template("""
+                int idx = i*2;
+                
+                T mesh_1_0 = triangles[idx*3];
+                T mesh_1_1 = triangles[idx*3+1];
+                T mesh_1_2 = triangles[idx*3+2];
+                T mesh_2_0 = triangles[idx*3+3];
+                T mesh_2_1 = triangles[idx*3+4];
+                T mesh_2_2 = triangles[idx*3+5];
+
+                U p01x = vertices[mesh_1_1*3] - vertices[mesh_1_0*3];
+                U p01y = vertices[mesh_1_1*3+1] - vertices[mesh_1_0*3+1];
+                U p01z = vertices[mesh_1_1*3+2] - vertices[mesh_1_0*3+2];
+                U p02x = vertices[mesh_1_2*3] - vertices[mesh_1_0*3];
+                U p02y = vertices[mesh_1_2*3+1] - vertices[mesh_1_0*3+1];
+                U p02z = vertices[mesh_1_2*3+2] - vertices[mesh_1_0*3+2];
+                cross(p01x, p01y, p01z, p02x, p02y, p02z, output[3*i], output[3*i+1],output[3*i+2]);
+                
+                p01x = vertices[mesh_2_1*3] - vertices[mesh_2_0*3];
+                p01y = vertices[mesh_2_1*3+1] - vertices[mesh_2_0*3+1];
+                p01z = vertices[mesh_2_1*3+2] - vertices[mesh_2_0*3+2];
+                p02x = vertices[mesh_2_2*3] - vertices[mesh_2_0*3];
+                p02y = vertices[mesh_2_2*3+1] - vertices[mesh_2_0*3+1];
+                p02z = vertices[mesh_2_2*3+2] - vertices[mesh_2_0*3+2];
+                cross(p01x, p01y, p01z, p02x, p02y, p02z, output[3*i], output[3*i+1],output[3*i+2]);
+                
+                output[3*i] /= 2.0;
+                output[3*i+1] /= 2.0;
+                output[3*i+2] /= 2.0;
+                
+                """).substitute(),
+            name="compute_meshes_normals_kernel",
+        )
+        return compute_meshes_normals_kernel
+
+
+    compute_meshes_normals_gpu_ = compute_meshes_normals_kernel()
+
+
+    def compute_meshes_normals_gpu(num_rows, num_cols, vertices, triangles):
+        vertices_gpu = cp.array(vertices, dtype=np.float32)
+        triangles_gpu = cp.array(triangles, dtype=np.int32)
+        output = cp.zeros((num_rows - 1, num_cols - 1, 3), dtype=np.float32)
+        compute_meshes_normals_gpu_(vertices_gpu,
+                                    triangles_gpu,
+                                    output,
+                                    size=((num_rows - 1) * (num_cols - 1)))
+        output_torch = torch.as_tensor(output)
+        return output_torch
+
+
+def compute_meshes_normals(num_rows, num_cols, vertices, triangles):
+    """
+    Compute normal vectors for all trimesh of the designed terrain
+
+    Parameters:
+        num_rows (int): Number of the row of the terrain
+        num_cols (int): Number of the column of the terrain
+        vertices (np.array(float)): array of shape (num_vertices, 3). Each row represents the location of each vertex [meters]
+        triangles (np.array(int)): array of shape (num_triangles, 3). Each row represents the indices of the 3 vertices connected by this triangle.
+    Returns:
+        normals (np.array(float)): array of shape (num_rows-1, num_cols-1, 3). Normals of the terrain
+    """
+    normals = np.zeros((num_rows - 1, num_cols - 1, 3), dtype=np.float32)
+    for i in range(num_rows - 1):
+        start = 2 * i * (num_cols - 1)
+        for j in range(num_cols - 1):
+            idx_1 = start + 2 * j
+            mesh_1, mesh_2 = triangles[idx_1], triangles[idx_1 + 1]
+            normal1 = compute_triangle_normal(vertices[mesh_1[0]],
+                                              vertices[mesh_1[1]],
+                                              vertices[mesh_1[2]])
+            normal2 = compute_triangle_normal(vertices[mesh_2[0]],
+                                              vertices[mesh_2[1]],
+                                              vertices[mesh_2[2]])
+            new_normal = (
+                                 normal1 +
+                                 normal2) / 2.0  # np.mean(np.stack((normal1, normal2)), axis=0)
+            if new_normal[2] < 0.:
+                new_normal *= -1
+            normals[i][j] = new_normal
+    return torch.from_numpy(normals)
+
+
+def compute_triangle_normal(p0, p1, p2):
+    """
+    Compute the normal vector of a trimesh
+
+    Parameters:
+        p0 (np.array(float)): array of shape (3,). Position of the trimesh's first vertice
+        p1 (np.array(float)): array of shape (3,). Position of the trimesh's second vertice
+        p2 (np.array(float)): array of shape (3,). Position of the trimesh's third vertice
+    Returns:
+        normal (np.array(float)): array of shape (3,). Normal Vector of the given trimesh
+    """
+    p0p1 = p1 - p0
+    p0p2 = p2 - p0
+    n = np.cross(p0p1, p0p2)
+    l = np.linalg.norm(n)
+    if l != 0.0:
+        normal = n / l
+    else:
+        normal = n
+    return normal
+
+
+class Point:
+
+    def __init__(self, x, y, z):
+        self.x = x
+        self.y = y
+        self.z = z
+
+
+@torch.jit.script
+def get_euler_xyz(q):
+    qx, qy, qz, qw = 0, 1, 2, 3
+    # roll (x-axis rotation)
+    sinr_cosp = 2.0 * (q[:, qw] * q[:, qx] + q[:, qy] * q[:, qz])
+    cosr_cosp = q[:, qw] * q[:, qw] - q[:, qx] * \
+                q[:, qx] - q[:, qy] * q[:, qy] + q[:, qz] * q[:, qz]
+    roll = torch.atan2(sinr_cosp, cosr_cosp)
+
+    # pitch (y-axis rotation)
+    sinp = 2.0 * (q[:, qw] * q[:, qy] - q[:, qz] * q[:, qx])
+    pitch = torch.where(
+        torch.abs(sinp) >= 1, copysign(np.pi / 2.0, sinp), torch.asin(sinp))
+
+    # yaw (z-axis rotation)
+    siny_cosp = 2.0 * (q[:, qw] * q[:, qz] + q[:, qx] * q[:, qy])
+    cosy_cosp = q[:, qw] * q[:, qw] + q[:, qx] * \
+                q[:, qx] - q[:, qy] * q[:, qy] - q[:, qz] * q[:, qz]
+    yaw = torch.atan2(siny_cosp, cosy_cosp)
+
+    return torch.stack((roll, pitch, yaw), dim=-1)
+
+
+@torch.jit.script
+def coordinate_rotation(axis: int, angle):
+    s = torch.sin(angle)
+    c = torch.cos(angle)
+
+    R = torch.eye(3, dtype=torch.float, device=angle.device)
+    R = R.reshape((1, 3, 3)).repeat(angle.size(0), 1, 1)
+
+    if axis == 0:
+        R[:, 1, 1] = c
+        R[:, 2, 2] = c
+        R[:, 1, 2] = s
+        R[:, 2, 1] = -s
+    elif axis == 1:
+        R[:, 0, 0] = c
+        R[:, 0, 2] = -s
+        R[:, 2, 0] = s
+        R[:, 2, 2] = c
+    elif axis == 2:
+        R[:, 0, 0] = c
+        R[:, 0, 1] = s
+        R[:, 1, 0] = -s
+        R[:, 1, 1] = c
+
+    return R
+
+
+@torch.jit.script
+def euler_xyz_to_rotation_matrix(rpy):
+    N = rpy.size(0)
+    R = torch.tensor((N, 3, 3), dtype=torch.float, device=rpy.device)
+    R = torch.matmul(
+        torch.matmul(coordinate_rotation(0, rpy[:, 0]),
+                     coordinate_rotation(1, rpy[:, 1])),
+        coordinate_rotation(2, rpy[:, 2]))
+    return R.transpose(1, 2)
+
+
+def get_contact_normals(robots_feet_pos, mesh_normals, contact_forces):
+    pos_to_idx = robots_feet_pos[:, :, :2].to(torch.int)
+
+    pos_to_idx = torch.permute(pos_to_idx, (2, 0, 1))
+    pos_to_idx = torch.flatten(pos_to_idx, 1)
+    contact_normals = mesh_normals[pos_to_idx.tolist()]
+    contact_normals = contact_normals.view(-1, 4, 3)
+
+    filter = torch.norm(contact_forces[:, :, :], dim=2) > 1.
+    filter = torch.repeat_interleave(filter, repeats=3,
+                                     dim=1).view(-1, robots_feet_pos.shape[1],
+                                                 3)
+    contact_normals *= filter
+
+    return contact_normals.flatten(1)
+
+
+@torch.jit.script
+def vector_2_skewmat(v):
+    N = v.size(0)
+    a = torch.zeros((N, 3, 3), device=v.device, dtype=torch.float)
+    a[:, 2, 1] = v[:, 0]
+    a[:, 2, 0] = -v[:, 1]
+    a[:, 1, 0] = v[:, 2]
+    a -= a.clone().transpose(1, 2)
+    return a
+
+
+def parse_sim_params(args, cfg):
+    if 'real' in args.task:
+        return cfg
+    from isaacgym import gymapi, gymutil
+    # initialize sim params
+    sim_params = gymapi.SimParams()
+    if 'flex' in args.physics_engine:
+        args.physics_engine = gymapi.SIM_FLEX
+        if args.device != "cpu":
+            print("WARNING: Using Flex with GPU instead of PHYSX!")
+    else:
+        args.physics_engine = gymapi.SIM_PHYSX
+        sim_params.physx.use_gpu = True if 'cuda' in args.sim_device else False
+        sim_params.physx.num_subscenes = args.subscenes
+    sim_params.use_gpu_pipeline = True if 'cuda' in args.rl_device else False
+    # if sim options are provided in cfg, parse them and update/override above:
+    if "sim" in cfg:
+        gymutil.parse_sim_config(cfg["sim"], sim_params)
+    # Override num_threads if passed on the command line
+    if args.physics_engine == gymapi.SIM_PHYSX and args.num_threads > 0:
+        sim_params.physx.num_threads = args.num_threads
+
+    return sim_params