diff --git a/legged_gym/envs/base/legged_robot.py b/legged_gym/envs/base/legged_robot.py index 2652f91..98ae04d 100644 --- a/legged_gym/envs/base/legged_robot.py +++ b/legged_gym/envs/base/legged_robot.py @@ -13,8 +13,8 @@ from typing import Tuple, Dict from legged_gym import LEGGED_GYM_ROOT_DIR from legged_gym.envs.base.base_task import BaseTask -from legged_gym.utils.math import quat_apply_yaw, wrap_to_pi, torch_rand_sqrt_float -from legged_gym.utils.isaacgym_utils import get_euler_xyz +from legged_gym.utils.math import wrap_to_pi +from legged_gym.utils.isaacgym_utils import get_euler_xyz as get_euler_xyz_in_tensor from legged_gym.utils.helpers import class_to_dict from .legged_robot_config import LeggedRobotCfg @@ -93,7 +93,7 @@ class LeggedRobot(BaseTask): # prepare quantities self.base_pos[:] = self.root_states[:, 0:3] self.base_quat[:] = self.root_states[:, 3:7] - self.rpy[:] = get_euler_xyz(self.base_quat) + self.rpy[:] = get_euler_xyz_in_tensor(self.base_quat[:]) self.base_lin_vel[:] = quat_rotate_inverse(self.base_quat, self.root_states[:, 7:10]) self.base_ang_vel[:] = quat_rotate_inverse(self.base_quat, self.root_states[:, 10:13]) self.projected_gravity[:] = quat_rotate_inverse(self.base_quat, self.gravity_vec) @@ -427,7 +427,7 @@ class LeggedRobot(BaseTask): self.dof_pos = self.dof_state.view(self.num_envs, self.num_dof, 2)[..., 0] self.dof_vel = self.dof_state.view(self.num_envs, self.num_dof, 2)[..., 1] self.base_quat = self.root_states[:, 3:7] - self.rpy = get_euler_xyz(self.base_quat) + self.rpy = get_euler_xyz_in_tensor(self.base_quat) self.base_pos = self.root_states[:self.num_envs, 0:3] self.contact_forces = gymtorch.wrap_tensor(net_contact_forces).view(self.num_envs, -1, 3) # shape: num_envs, num_bodies, xyz axis diff --git a/legged_gym/envs/base/legged_robot_config.py b/legged_gym/envs/base/legged_robot_config.py index f029268..03f09a2 100644 --- a/legged_gym/envs/base/legged_robot_config.py +++ b/legged_gym/envs/base/legged_robot_config.py @@ -9,6 +9,7 @@ class LeggedRobotCfg(BaseConfig): env_spacing = 3. # not used with heightfields/trimeshes send_timeouts = True # send time out information to the algorithm episode_length_s = 20 # episode length in seconds + test = False class terrain: mesh_type = 'plane' # "heightfield" # none, plane, heightfield or trimesh diff --git a/legged_gym/envs/h1/h1_config.py b/legged_gym/envs/h1/h1_config.py index 64802b1..ba4c3a1 100644 --- a/legged_gym/envs/h1/h1_config.py +++ b/legged_gym/envs/h1/h1_config.py @@ -28,7 +28,6 @@ class H1RoughCfg( LeggedRobotCfg ): class env(LeggedRobotCfg.env): num_observations = 42 num_actions = 10 - test = False class control( LeggedRobotCfg.control ): diff --git a/legged_gym/utils/isaacgym_utils.py b/legged_gym/utils/isaacgym_utils.py index 3fd8c0c..a92c0df 100644 --- a/legged_gym/utils/isaacgym_utils.py +++ b/legged_gym/utils/isaacgym_utils.py @@ -3,474 +3,11 @@ 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) @@ -491,94 +28,3 @@ def get_euler_xyz(q): 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