Isaac Gym, UR5 Inverse Kinematics to target, CPU vs GPU differences

UR5_IK.py

"""
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)

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:

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

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.

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

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.