585 lines
19 KiB
Python
585 lines
19 KiB
Python
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
|