update
This commit is contained in:
brook-bee
parent
561e121a42
commit
d816fe187a
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -28,7 +28,6 @@ class H1RoughCfg( LeggedRobotCfg ):
|
|||
class env(LeggedRobotCfg.env):
|
||||
num_observations = 42
|
||||
num_actions = 10
|
||||
test = False
|
||||
|
||||
|
||||
class control( LeggedRobotCfg.control ):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Reference in New Issue