def push_in_dir(self, dir, amount, T, stiffness=1200):
delta_t = 0.1
for env_index, env_ptr in enumerate(self._scene.env_ptrs):
block_ah = self._scene.ah_map[env_index][self._block_name]
block_transform = self._block.get_rb_transforms(env_ptr,
block_ah)[0]
block_start = transform_to_np(block_transform)
ee_start = self._frankas[env_index].get_ee_transform(
env_ptr, self._franka_name)
np_ee_start = transform_to_np(ee_start)
des_quat = np_to_quat(np_ee_start[3:], format="wxyz")
block_desired = np.array([
block_start[0] - amount * np.sin(self.dir_to_rpy[dir][2]),
block_start[1],
block_start[2] - amount * np.cos(self.dir_to_rpy[dir][2])
])
t = 0
self._frankas[env_index].set_attractor_props(
env_index, env_ptr, self._franka_name, {
'stiffness': stiffness,
'damping': 4 * np.sqrt(stiffness)
})
goal_tol = 0.02
while (t < T):
ee_pose = self._frankas[env_index].get_ee_transform(
env_ptr, self._franka_name)
np_pose = transform_to_np(ee_pose)[0:3]
block_transform_np = transform_to_np(
self._block.get_rb_transforms(env_ptr, block_ah)[0])
distance = np.linalg.norm(block_transform_np[:3] -
block_desired)
#print("Distance", distance)
des_ee_trans = np.array([
np_pose[0] - distance * np.sin(self.dir_to_rpy[dir][2]),
np_pose[1],
np_pose[2] - distance * np.cos(self.dir_to_rpy[dir][2])
])
#np_pose[2] += action[0]
transformed_pt = gymapi.Transform(p=np_to_vec3(des_ee_trans),
r=ee_start.r)
self._frankas[env_index].set_attractor_props(
env_index, env_ptr, self._franka_name, {
'stiffness': stiffness,
'damping': 4 * np.sqrt(stiffness)
})
self._frankas[env_index].set_ee_transform(
env_ptr, env_index, self._franka_name, transformed_pt)
time_taken = self.goto_pose(des_quat,
des_ee_trans,
env_index,
env_ptr,
transformed_pt,
max_t=T // 5)
t += time_taken
if distance < goal_tol:
break
def goto_side(self, dir, stiffness=1000):
"""
Goes to the direction in "0,1,2,3"
0
3 1
2
"""
#base is -1.57 1.57 1.56
des_quat = rpy_to_quat(np.array(self.dir_to_rpy[dir]))
delta_side = 0.02 + self._cfg["block"]["dims"]["width"] / 2
up_offset = self._cfg["block"]["dims"]["height"] / 2 + 0.07
for env_index, env_ptr in enumerate(self._scene.env_ptrs):
block_ah = self._scene.ah_map[env_index][self._block_name]
block_transform = self._block.get_rb_transforms(env_ptr,
block_ah)[0]
block_transform_trans = vec3_to_np(block_transform.p)
des_ee_trans = np.array([
block_transform_trans[0] +
delta_side * np.sin(self.dir_to_rpy[dir][2]),
block_transform_trans[1], block_transform_trans[2] +
delta_side * np.cos(self.dir_to_rpy[dir][2])
])
up_des_ee_trans = des_ee_trans.copy()
up_des_ee_trans[1] += up_offset
#yzx
#IK to compute joint angles
transformed_pt = gymapi.Transform(p=np_to_vec3(des_ee_trans),
r=des_quat)
up_transformed_pt = gymapi.Transform(p=np_to_vec3(up_des_ee_trans),
r=des_quat)
self._frankas[env_index].set_attractor_props(
env_index, env_ptr, self._franka_name, {
'stiffness': stiffness,
'damping': 4 * np.sqrt(stiffness)
})
ee_pose = self._frankas[env_index].get_ee_transform(
env_ptr, self._franka_name)
if ee_pose.p.y < transformed_pt.p.y + up_offset:
np_ee_pose = transform_to_np(ee_pose, format="wxyz")
np_ee_pose[1] = transformed_pt.p.y + up_offset
ee_pose.p.y = transformed_pt.p.y + up_offset
self.goto_pose(np_to_quat(np_ee_pose[3:], format="wxyz"),
np_ee_pose[0:3], env_index, env_ptr, ee_pose)
self.goto_pose(des_quat, up_des_ee_trans, env_index, env_ptr,
up_transformed_pt)
self.goto_pose(des_quat, des_ee_trans, env_index, env_ptr,
transformed_pt)
self._frankas[env_index].set_attractor_props(
env_index, env_ptr, self._franka_name, {
'stiffness': 500,
'damping': 4 * np.sqrt(stiffness)
})
def get_block_poses(self, env_idx=None):
if env_idx is None:
box_pose_obs = np.zeros((self.n_envs, 7))
for env_index, env_ptr in enumerate(self._scene.env_ptrs):
ah = self._scene.ah_map[env_index][self._block_name]
block_transform = self._block.get_rb_transforms(env_ptr, ah)[0]
block_transform_np = transform_to_np(block_transform,
format="wxyz")
if env_idx == env_index:
return block_transform_np
else:
box_pose_obs[env_index, :] = block_transform_np
return box_pose_obs
def _update_state(self):
robot_pos_fqn = "frame:pose/position"
robot_orn_fqn = "frame:pose/quaternion"
block_pos_fqn = "frame:block:pose/position"
block_orn_fqn = "frame:block:pose/quaternion"
obs_pos_fqn = "frame:obstacle:pose/position"
obs_orn_fqn = "frame:obstacle:pose/quaternion"
block_on_color_fqn = "frame:block:on_color"
for env_index, env_ptr in enumerate(self._scene.env_ptrs):
block_ah = self._scene.ah_map[env_index][self._block_name]
obs_ah = self._scene.ah_map[env_index][self._obstacle_name]
block_transform_np = transform_to_np(self._block.get_rb_transforms(
env_ptr, block_ah)[0],
format="wxyz")
obs_transform_np = transform_to_np(self._block.get_rb_transforms(
env_ptr, obs_ah)[0],
format="wxyz")
ee_pose_np = transform_to_np(
self._frankas[env_index].get_ee_transform(
env_ptr, self._franka_name),
format="wxyz")
#all_cts = self._scene.gym.get_rigid_contacts(self._scene._sim) # check for between block and board
#all_cts = all_cts[all_cts['env0'] != -1 & np.logical_not(np.isclose(all_cts['lambda'], 0))]
self.pillar_states[env_index].set_values_from_vec(
[robot_pos_fqn], ee_pose_np[:3].tolist())
self.pillar_states[env_index].set_values_from_vec(
[robot_orn_fqn], ee_pose_np[3:].tolist())
self.pillar_states[env_index].set_values_from_vec(
[block_pos_fqn], block_transform_np[:3].tolist())
self.pillar_states[env_index].set_values_from_vec(
[block_orn_fqn], block_transform_np[3:].tolist())
self.pillar_states[env_index].set_values_from_vec(
[obs_pos_fqn], obs_transform_np[:3].tolist())
self.pillar_states[env_index].set_values_from_vec(
[obs_orn_fqn], obs_transform_np[3:].tolist())
color = color_block_is_on(self._cfg, block_transform_np)
self.pillar_states[env_index].set_values_from_vec(
[block_on_color_fqn], color)
def _apply_actions(self, action, planning_env=False):
for env_index, env_ptr in enumerate(self._scene.env_ptrs):
ee_pose = self._frankas[env_index].get_ee_transform(
env_ptr, self._franka_name)
np_pose = transform_to_np(ee_pose)[0:3]
np_pose[2] += action[0]
transformed_pt = gymapi.Transform(p=np_to_vec3(np.array(np_pose)),
r=ee_pose.r)
stiffness = action[1]
self._frankas[env_index].set_attractor_props(
env_index, env_ptr, self._franka_name, {
'stiffness': stiffness,
'damping': 4 * np.sqrt(stiffness)
})
self._frankas[env_index].set_ee_transform(env_ptr, env_index,
self._franka_name,
transformed_pt)
def goto_pose(self,
des_quat,
des_ee_trans,
env_index,
env_ptr,
transformed_pt,
max_t=None):
pos_tol = 0.005
quat_tol = 1.5
self._frankas[env_index].set_ee_transform(env_ptr, env_index,
self._franka_name,
transformed_pt)
if max_t is None:
max_t = self.max_steps_per_movement
for i in range(int(max_t)):
self._scene.step()
if i % 10:
self.render(custom_draws=custom_draws)
if i % 20:
ee_pose = self._frankas[env_index].get_ee_transform(
env_ptr, self._franka_name)
np_pose = transform_to_np(ee_pose, format="wxyz")
if np.linalg.norm(des_ee_trans - np_pose[:3]
) < pos_tol and Quaternion.absolute_distance(
Quaternion(
quat_to_np(des_quat, format="wxyz")),
Quaternion(np_pose[3:])) < quat_tol:
[(self._scene.step(),
self.render(custom_draws=custom_draws))
for i in range(10)]
break
if i == max_t - 1:
print(
"COuld not reach pose in time, distance still {} after time {}"
.format(np.linalg.norm(des_ee_trans - np_pose[:3]), i))
return i