Skip to content

Instantly share code, notes, and snippets.

@DanielTakeshi
Last active January 31, 2024 16:37
Show Gist options
  • Star 1 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save DanielTakeshi/84d6e83c5f7ea89b5f436f60d99fb610 to your computer and use it in GitHub Desktop.
Save DanielTakeshi/84d6e83c5f7ea89b5f436f60d99fb610 to your computer and use it in GitHub Desktop.
Isaac Gym, UR5 Inverse Kinematics to target, CPU vs GPU differences
"""
Runs IK to get the UR5 end-effector to reach a target. Inspect CPU vs GPU mode.
Use a yellow sphere to show the tip of the end-effector, and a blue sphere to
show the target. These spheres are only used for debugging / visualization.
"""
from isaacgym import gymapi
from isaacgym import gymutil
from isaacgym import gymtorch
from isaacgym.torch_utils import (quat_conjugate, quat_mul, quat_apply)
import numpy as np
import torch
torch.set_printoptions(precision=4, sci_mode=False, linewidth=120)
gym_GREEN = gymapi.Vec3(0., 1., 0.)
gym_BLUE = gymapi.Vec3(0., 0., 1.)
# ---------------------------------------------------------------------------- #
# Helper methods.
# ---------------------------------------------------------------------------- #
def control_ik(j_eef, dpose, num_envs, num_dofs, damping=0.05):
"""Solve damped least squares, from `franka_cube_ik_osc.py` in Isaac Gym.
Returns: Change in DOF positions, [num_envs,num_dofs], to add to current positions.
"""
j_eef_T = torch.transpose(j_eef, 1, 2)
lmbda = torch.eye(6).to(j_eef_T.device) * (damping ** 2)
u = (j_eef_T @ torch.inverse(j_eef @ j_eef_T + lmbda) @ dpose).view(num_envs, num_dofs)
return u
def orientation_error(desired, current):
cc = quat_conjugate(current)
q_r = quat_mul(desired, cc)
return q_r[:, 0:3] * torch.sign(q_r[:, 3]).unsqueeze(-1)
def move_sphere_ee_vectorized(hand_pos, hand_rot, eetip_offsets):
"""Move spheres representing the EE tip by applying an offset correction.
Returns: eetip_pos, shape (num_envs,3)
"""
eetip_pos = hand_pos + quat_apply(hand_rot, eetip_offsets)
return eetip_pos
def get_target_orientation(direction='top_down'):
"""Some sample target poses."""
if direction == 'top_down':
return gymapi.Quat(1.0, 0.0, 0.0, 0.0)
elif direction == 'bottom_up':
return gymapi.Quat(0.0, 0.0, 0.0, 1.0)
else:
raise ValueError(direction)
def sample_sphere_surface(center_sphere, radius, n_points=1):
"""Sample about a sphere surface, centered at `center_sphere`.
Samples IID standard Gaussians. Then normalize and multiply each by the radius.
Returns: points, shaped (N,3).
"""
assert radius > 0
xyz_N3 = np.random.normal(loc=0.0, scale=1.0, size=(n_points, 3))
xyz_N3 = xyz_N3 * (radius / np.linalg.norm(xyz_N3, axis=1, keepdims=True))
return center_sphere + xyz_N3
def get_sim_params(args):
"""Start up a common set of simulation parameters.
Based on `franka_cube_ik_osc.py` provided by Isaac Gym authors.
"""
sim_params = gymapi.SimParams()
sim_params.up_axis = gymapi.UP_AXIS_Z
sim_params.gravity = gymapi.Vec3(0.0, 0.0, -9.8)
sim_params.dt = 1.0 / 60.0
sim_params.substeps = 2
sim_params.use_gpu_pipeline = args.use_gpu_pipeline
if args.physics_engine == gymapi.SIM_PHYSX:
sim_params.physx.solver_type = 1
sim_params.physx.num_position_iterations = 8
sim_params.physx.num_velocity_iterations = 1
sim_params.physx.rest_offset = 0.0
sim_params.physx.contact_offset = 0.001
sim_params.physx.friction_offset_threshold = 0.001
sim_params.physx.friction_correlation_distance = 0.0005
sim_params.physx.num_threads = args.num_threads
sim_params.physx.use_gpu = args.use_gpu
else:
raise Exception("This example can only be used with PhysX")
return sim_params
def get_UR5_asset(gym, sim, asset_root, asset_file):
"""Create a UR5 asset with a linear slider."""
asset_options = gymapi.AssetOptions()
asset_options.armature = 0.01
asset_options.fix_base_link = True
asset_options.disable_gravity = True
asset_options.flip_visual_attachments = True
ur5_asset = gym.load_asset(sim, asset_root, asset_file, asset_options)
return ur5_asset
def get_sphere_asset(gym, sim, radius):
"""Create a sphere asset and disable gravity."""
asset_options = gymapi.AssetOptions()
asset_options.disable_gravity = True
sphere_asset = gym.create_sphere(sim, radius, asset_options)
return sphere_asset
# ---------------------------------------------------------------------------- #
# Acquire gym and set arguments. Then set up the envs, etc.
# ---------------------------------------------------------------------------- #
gym = gymapi.acquire_gym()
args = gymutil.parse_arguments(
description="IK example, debug CPU vs GPU differences.",
custom_parameters=[
{"name": "--num_envs", "type": int, "default": 16},
{"name": "--seed", "type": int, "default": 10},
]
)
# Bells and whistles.
np.random.seed(args.seed)
torch.manual_seed(args.seed)
num_envs = args.num_envs
args.use_gpu = args.use_gpu_pipeline
sim_params = get_sim_params(args)
device = args.sim_device if args.use_gpu_pipeline else 'cpu'
print(f"args: {args}\ndevice: {device}")
# Create sim
sim = gym.create_sim(
args.compute_device_id, args.graphics_device_id, args.physics_engine, sim_params
)
# Create viewer.
camera_props = gymapi.CameraProperties()
camera_props.horizontal_fov = 75.0
camera_props.width = 1920
camera_props.height = 1080
viewer = gym.create_viewer(sim, camera_props)
# Add ground plane
plane_params = gymapi.PlaneParams()
plane_params.normal = gymapi.Vec3(0, 0, 1)
gym.add_ground(sim, plane_params)
# Get UR5, target, and EE tip assets.
asset_root = "assets"
asset_file = "ur5.urdf"
ur5_asset = get_UR5_asset(gym, sim, asset_root, asset_file)
targ_asset = get_sphere_asset(gym, sim, radius=0.03)
eetip_asset = get_sphere_asset(gym, sim, radius=0.03)
# Configure UR5 dofs, following Franka example.
ur5_dof_props = gym.get_asset_dof_properties(ur5_asset)
ur5_lower_limits = ur5_dof_props["lower"]
ur5_upper_limits = ur5_dof_props["upper"]
ur5_ranges = ur5_upper_limits - ur5_lower_limits
ur5_mids = 0.5 * (ur5_upper_limits + ur5_lower_limits)
# UR5-specific starting pose.
ur5_pose = gymapi.Transform()
ur5_position = [0.0, 0.0, 1.453]
ur5_pose.p = gymapi.Vec3(ur5_position[0], ur5_position[1], ur5_position[2])
ur5_pose.r = gymapi.Quat(0.0, 0.707107, 0.0, 0.707107)
# Use position drive for all dofs, following Franka example.
ur5_dof_props["driveMode"].fill(gymapi.DOF_MODE_POS)
ur5_dof_props["stiffness"][:8].fill(400.0)
ur5_dof_props["damping"][:8].fill(40.0)
# Default dof states and position targets, following Franka example.
ur5_num_dofs = gym.get_asset_dof_count(ur5_asset)
default_dof_pos = np.zeros(ur5_num_dofs, dtype=np.float32)
default_dof_pos[:7] = ur5_mids[:7]
default_dof_state = np.zeros(ur5_num_dofs, gymapi.DofState.dtype)
default_dof_state["pos"] = default_dof_pos
# Get link index for end-effector.
ur5_link_dict = gym.get_asset_rigid_body_dict(ur5_asset)
ur5_ee_index = ur5_link_dict["tool0"]
# Set up the grid of environments
num_per_row = int(np.sqrt(num_envs))
spacing = 1.5
env_lower = gymapi.Vec3(-spacing, -spacing, 0.0)
env_upper = gymapi.Vec3(spacing, spacing, spacing)
# Information to cache.
envs = []
targ_handles = []
targ_idxs = []
ee_idxs = []
eetip_handles = []
eetip_idxs = []
init_pos_list = []
init_rot_list = []
# Use sphere sampling to sample targets for IK.
target_points = sample_sphere_surface(
center_sphere=(1.0, 0.0, 1.9), radius=0.01, n_points=num_envs,
)
for i in range(num_envs):
env = gym.create_env(sim, env_lower, env_upper, num_per_row)
envs.append(env)
# Actor 0: UR5, set dof properties and initial DOF states / targets.
ur5_handle = gym.create_actor(env, ur5_asset, ur5_pose, "ur5", i, 0)
gym.set_actor_dof_properties(env, ur5_handle, ur5_dof_props)
gym.set_actor_dof_states(env, ur5_handle, default_dof_state, gymapi.STATE_ALL)
gym.set_actor_dof_position_targets(env, ur5_handle, default_dof_pos)
# Actor 1: Create targets for IK. No collisions with other objects. Also, set the
# orientation here which we can query from later to get a desired angle for IK.
targ_pose = gymapi.Transform()
targ_pose.p = gymapi.Vec3(*target_points[i,:])
targ_pose.r = get_target_orientation(direction='bottom_up')
targ_id = f"targ_{i}"
targ_coll = num_envs + 1
targ_handle = gym.create_actor(env, targ_asset, targ_pose, targ_id, targ_coll, 0)
gym.set_rigid_body_color(env, targ_handle, 0, gymapi.MESH_VISUAL, gym_BLUE)
targ_idx = gym.get_actor_rigid_body_index(env, targ_handle, 0, gymapi.DOMAIN_SIM)
targ_handles.append(targ_handle)
targ_idxs.append(targ_idx)
# Actor 2: debugging spheres for the EE tip since we don't have it in the URDF.
eetip_pose = gymapi.Transform()
eetip_id = f"eetip_{i}"
eetip_coll = num_envs + 2
eetip_handle = gym.create_actor(env, eetip_asset, eetip_pose, eetip_id, eetip_coll, 0)
gym.set_rigid_body_color(env, eetip_handle, 0, gymapi.MESH_VISUAL, gym_GREEN)
eetip_idx = gym.get_actor_rigid_body_index(env, eetip_handle, 0, gymapi.DOMAIN_SIM)
eetip_handles.append(eetip_handle)
eetip_idxs.append(eetip_idx)
# Track the indices of the 'last' UR5 link. For the 'real tip' see `eetip_idx`.
ee_idx = gym.find_actor_rigid_body_index(env, ur5_handle, "tool0", gymapi.DOMAIN_SIM)
ee_idxs.append(ee_idx)
# Get inital hand pose, this is the same across all envs.
ee_handle = gym.find_actor_rigid_body_handle(env, ur5_handle, "tool0")
ee_pose = gym.get_rigid_transform(env, ee_handle)
init_pos_list.append([ee_pose.p.x, ee_pose.p.y, ee_pose.p.z])
init_rot_list.append([ee_pose.r.x, ee_pose.r.y, ee_pose.r.z, ee_pose.r.w])
# point camera at middle env
cam_pos = gymapi.Vec3(4, 3, 2)
cam_target = gymapi.Vec3(-4, -3, 0)
middle_env = envs[args.num_envs // 2 + num_per_row // 2]
gym.viewer_camera_look_at(viewer, middle_env, cam_pos, cam_target)
# ============================== prepare tensors ===================================
# from now on, we will use the tensor API that can run on CPU or GPU
gym.prepare_sim(sim)
# initial hand position and orientation tensors
init_pos = torch.Tensor(init_pos_list).view(num_envs, 3).to(device)
init_rot = torch.Tensor(init_rot_list).view(num_envs, 4).to(device)
# eetip offsets for us to track [to debug: turn off robot movement]
eetip_offsets = torch.tensor([0., 0., 0.18], device=device).repeat((num_envs,1))
# get jacobian tensor, I think tensor shape (num_envs, num_links, [pose], num_joints-dof)
_jacobian = gym.acquire_jacobian_tensor(sim, "ur5")
jacobian = gymtorch.wrap_tensor(_jacobian)
# jacobian entries corresponding to ur5 ee
j_eef = jacobian[:, ur5_ee_index - 1, :, :]
# Actor root state tensor, only to control debugging sphere target if desired.
_actor_root_state_tensor = gym.acquire_actor_root_state_tensor(sim)
root_state_tensor = gymtorch.wrap_tensor(_actor_root_state_tensor).view(num_envs, -1, 13)
# Get rigid body state tensor, shape (num_RBs, 13).
_rb_states = gym.acquire_rigid_body_state_tensor(sim)
rb_states = gymtorch.wrap_tensor(_rb_states)
# Get dof state tensor, we will be querying to get current DOF state, and adding to it.
_dof_states = gym.acquire_dof_state_tensor(sim)
dof_states = gymtorch.wrap_tensor(_dof_states)
ur5_dof_states = dof_states.view(num_envs, -1, 2)[:,:ur5_num_dofs]
dof_pos = ur5_dof_states[:,:,0]
dof_pos = torch.unsqueeze(dof_pos, 2).to(device)
total_dofs = gym.get_sim_dof_count(sim) // num_envs
while not gym.query_viewer_has_closed(viewer):
ur5_dof_targets = torch.zeros((num_envs, total_dofs), dtype=torch.float, device=device)
# Step the physics
gym.simulate(sim)
gym.fetch_results(sim, True)
# Refresh Tensors
gym.refresh_actor_root_state_tensor(sim)
gym.refresh_rigid_body_state_tensor(sim)
gym.refresh_dof_state_tensor(sim)
gym.refresh_jacobian_tensors(sim)
# Due to refreshing tensors, `rb_states` gets updated information.
goal_pos = rb_states[targ_idxs, :3]
goal_rot = rb_states[targ_idxs, 3:7] # we should assign to this beforehand
hand_pos = rb_states[ee_idxs, :3] # this is the UR5 joint 'before' the tip
hand_rot = rb_states[ee_idxs, 3:7]
# Adjust the 'current position' to be the EE tip.
curr_pos = move_sphere_ee_vectorized(hand_pos, hand_rot, eetip_offsets)
# Compute position and (if desired) orientation error to the goal.
pos_err = goal_pos - curr_pos # (num_envs,3)
orn_err = orientation_error(goal_rot, hand_rot) # (num_envs,3)
dpose = torch.cat([pos_err, orn_err], -1).unsqueeze(-1) # (num_envs,6,1)
pos_err_norm = torch.norm(pos_err, dim=-1).unsqueeze(-1) # can print if desired
orn_err_norm = torch.norm(orn_err, dim=-1).unsqueeze(-1) # can print if desired
# Use IK controller using damped least squares to update position targets.
u_delta = control_ik(j_eef, dpose, num_envs, num_dofs=ur5_num_dofs)
_pos_target = dof_pos.squeeze(-1) + u_delta # (num_envs,7)
_pos_target = _pos_target.view(num_envs, -1) # (num_envs,7)
ur5_dof_targets[:,:ur5_num_dofs] = _pos_target
# Move the robot by setting DOF position targets.
gym.set_dof_position_target_tensor(
sim, gymtorch.unwrap_tensor(ur5_dof_targets)
)
# Move the debugging sphere which represents the robot EE tip.
eetip_pos = move_sphere_ee_vectorized(hand_pos, hand_rot, eetip_offsets)
root_state_tensor[:, 2, 0:3] = eetip_pos # (num_envs,3), index 2 for sphere
gym.set_actor_root_state_tensor(
sim, gymtorch.unwrap_tensor(root_state_tensor)
)
# Update viewer
gym.step_graphics(sim)
gym.draw_viewer(viewer, sim, False)
gym.sync_frame_time(sim)
# cleanup
gym.destroy_viewer(viewer)
gym.destroy_sim(sim)
@DanielTakeshi
Copy link
Author

DanielTakeshi commented Nov 28, 2022

Instructions

I'm using Ubuntu 18.04 with an NVIDIA 3090 GPU.

Installation

Create a conda environment following the Isaac Gym installation instructions. (I'm using version 1.0rc4 for isaacgym.)

Get python and data files

To get all of the data files and this python script use this zip file: https://drive.google.com/file/d/1AJEzef4ufNRYvH2oXszhcUKWqy3aqEWb/view?usp=sharing This will contain everything you need, including this script (UR5_IK.py). Unzip it and you should see:

(rlgpu) seita@machine:~/MWE_UR5_IK$ ls -lh 
total 20K
drwxrwxr-x 3 seita seita 4.0K Nov 28 14:06 assets
-rw-rw-r-- 1 seita seita  14K Nov 28 14:06 UR5_IK.py
(rlgpu) seita@machine:~/MWE_UR5_IK$ ls -lh assets/
total 16K
drwxrwxr-x 2 seita seita 4.0K Nov 28 14:06 meshes
-rw-rw-r-- 1 seita seita  12K Nov 28 14:06 ur5.urdf
(rlgpu) seita@machine:~/MWE_UR5_IK$ ls -lh assets/meshes/
total 8.8M
-rw-rw-r-- 1 seita seita 153K Nov 28 14:06 base.dae
-rw-rw-r-- 1 seita seita  29K Nov 28 14:06 base.stl
-rw-rw-r-- 1 seita seita 2.1M Nov 28 14:06 ee3_coarse_edited_.stl
-rw-rw-r-- 1 seita seita 1.5M Nov 28 14:06 forearm.dae
-rw-rw-r-- 1 seita seita  52K Nov 28 14:06 forearm.stl
-rw-rw-r-- 1 seita seita 988K Nov 28 14:06 shoulder.dae
-rw-rw-r-- 1 seita seita  33K Nov 28 14:06 shoulder.stl
-rw-rw-r-- 1 seita seita 2.0M Nov 28 14:06 upperarm.dae
-rw-rw-r-- 1 seita seita  58K Nov 28 14:06 upperarm.stl
-rw-rw-r-- 1 seita seita 931K Nov 28 14:06 wrist1.dae
-rw-rw-r-- 1 seita seita  35K Nov 28 14:06 wrist1.stl
-rw-rw-r-- 1 seita seita 929K Nov 28 14:06 wrist2.dae
-rw-rw-r-- 1 seita seita  35K Nov 28 14:06 wrist2.stl
-rw-rw-r-- 1 seita seita 125K Nov 28 14:06 wrist3.dae
-rw-rw-r-- 1 seita seita  22K Nov 28 14:06 wrist3.stl
(rlgpu) seita@machine:~/MWE_UR5_IK$ 

Test CPU vs GPU mode

Then test CPU and GPU mode:

python UR5_IK.py --pipeline=cpu
python UR5_IK.py --pipeline=gpu

Observe that the results are different.

On the CPU the yellow/blue target spheres almost coincide (good) and this is what usually happens:

Screenshot from 2022-11-28 13-54-49

On the GPU, I consistently see a different steady state where the robot isn't able to get the two spheres to coincide:

Screenshot from 2022-11-28 13-56-53

Why does this happen? I am not sure.

The confusing part is, for the GPU mode, the u_delta we compute for the change in the target DOFs will show nonzero values, so IK is "trying" to get new target DOFs, but the robot isn't able to move.

Screenshot from 2022-11-28 14-14-57

Of course the bigger confusion is the CPU vs GPU mode discrepancy -- I thought they would give the same performance?

@RyanPaulMcKenna
Copy link

My only guess is that perhaps one of the torch functions or the isaac gym functions in torch utils behaves differently between cpu and gpu which would be a bug if that is the case.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment