170 lines
5.7 KiB
Python
170 lines
5.7 KiB
Python
import numpy as np
|
|
import time
|
|
from roboimi.envs.double_base import DualDianaMed
|
|
# import cv2
|
|
import roboimi.utils.KDL_utils.transform as T
|
|
|
|
import copy as cp
|
|
|
|
|
|
class DualDianaMed_Pos_Ctrl(DualDianaMed):
|
|
"""
|
|
|
|
:param is_render: Choose if use the renderer to render the scene or not.
|
|
:param renderer: Choose official renderer with "viewer",
|
|
another renderer with "mujoco_viewer"
|
|
:param control_freq: Upper-layer control frequency.
|
|
Note that high frequency will cause high time-lag.
|
|
:param cam_view: Choice of Camera view.
|
|
"""
|
|
|
|
def __init__(self,
|
|
robot=None,
|
|
is_render=False,
|
|
renderer="viewer",
|
|
control_freq=200,
|
|
is_interpolate=True,
|
|
cam_view=None
|
|
):
|
|
super().__init__(
|
|
robot=robot,
|
|
is_render=is_render,
|
|
renderer=renderer,
|
|
control_freq=control_freq,
|
|
is_interpolate=is_interpolate,
|
|
cam_view=cam_view
|
|
)
|
|
|
|
self.max_reward = 4
|
|
|
|
self.cam_start()
|
|
|
|
|
|
def step(self,action=np.zeros(16)):
|
|
action_left = self.ik_solve(action[:3],action[3:7],self.arm_left)
|
|
action_right = self.ik_solve(action[7:10],action[10:14],self.arm_right)
|
|
action = np.hstack((action_left,action_right,action[14:]))
|
|
super().step(action)
|
|
self.rew = self._get_reward()
|
|
|
|
def step_jnt(self,action):
|
|
super().step(action)
|
|
|
|
def ik_solve(self,pos_goal,quat_goal,Arm):
|
|
pos_bias,mat_bias = Arm.get_Arm_bias_frame()
|
|
pos_bias = np.array([pos_bias[0],pos_bias[1],pos_bias[2]])
|
|
mat_bias = np.array([[mat_bias[0],mat_bias[1],mat_bias[2]],
|
|
[mat_bias[3],mat_bias[4],mat_bias[5]],
|
|
[mat_bias[6],mat_bias[7],mat_bias[8]]])
|
|
mat_bias = np.linalg.inv(mat_bias)
|
|
mat_goal = T.quat2Mat(quat_goal)
|
|
p_goal = mat_bias@(pos_goal-pos_bias)
|
|
mat_goal = mat_bias@mat_goal
|
|
return Arm.kdl_solver.ikSolver(p_goal, mat_goal, Arm.arm_qpos)
|
|
|
|
def reset(self,box_pos):
|
|
|
|
self.mj_data.joint('red_box_joint').qpos[0] = box_pos[0]
|
|
self.mj_data.joint('red_box_joint').qpos[1] = box_pos[1]
|
|
self.mj_data.joint('red_box_joint').qpos[2] = box_pos[2]
|
|
self.mj_data.joint('red_box_joint').qpos[3] = 1.0
|
|
self.mj_data.joint('red_box_joint').qpos[4] = 0.0
|
|
self.mj_data.joint('red_box_joint').qpos[5] = 0.0
|
|
self.mj_data.joint('red_box_joint').qpos[6] = 0.0
|
|
super().reset()
|
|
self.top = None
|
|
self.angle = None
|
|
self.r_vis = None
|
|
self.front = None
|
|
self.cam_flage = True
|
|
t=0
|
|
while self.cam_flage:
|
|
if(type(self.top)==type(None)
|
|
or type(self.angle)==type(None)
|
|
or type(self.r_vis)==type(None)
|
|
or type(self.front)==type(None)):
|
|
time.sleep(0.001)
|
|
t+=1
|
|
else:
|
|
self.cam_flage=False
|
|
|
|
|
|
|
|
def preStep(self, action):
|
|
if isinstance(action,np.ndarray) and len(action)==16:
|
|
super().preStep(action)
|
|
else:
|
|
raise AttributeError("size should be 16")
|
|
|
|
def get_env_state(self):
|
|
box_pose = np.zeros(3)
|
|
for i in range(3):
|
|
box_pose[i] = cp.deepcopy(self.mj_data.joint('red_box_joint').qpos[i])
|
|
return box_pose
|
|
|
|
|
|
def _get_reward(self):
|
|
all_contact_pairs = []
|
|
for collision_num in range(self.mj_data.ncon):
|
|
geom1 = self.mj_data.contact[collision_num].geom1
|
|
geom2 = self.mj_data.contact[collision_num].geom2
|
|
name_geom_1 = self.getID2Name('geom',geom1)
|
|
name_geom_2 = self.getID2Name('geom',geom2)
|
|
contact_pair = (name_geom_1, name_geom_2)
|
|
all_contact_pairs.append(contact_pair)
|
|
|
|
touch_table = ("l_fingertip_g0_right", "table") in all_contact_pairs or\
|
|
("l_fingertip_g0_right", "table") in all_contact_pairs
|
|
box_touch_table = ("red_box", "table") in all_contact_pairs
|
|
touch_left_gripper = ("red_box", "r_finger_left") in all_contact_pairs or\
|
|
("red_box", "l_finger_left") in all_contact_pairs
|
|
touch_right_gripper = ("red_box", "r_finger_right") in all_contact_pairs or\
|
|
("red_box", "l_finger_right") in all_contact_pairs
|
|
|
|
reward = 0
|
|
if touch_right_gripper and not touch_table:
|
|
reward = 1
|
|
if touch_right_gripper and not box_touch_table:
|
|
reward = 2
|
|
if touch_left_gripper: # attempted transfer
|
|
reward = 3
|
|
if touch_left_gripper and not box_touch_table: # successful transfer
|
|
reward = 4
|
|
return reward
|
|
|
|
|
|
def make_sim_env(task_name):
|
|
if 'sim_transfer' in task_name:
|
|
from roboimi.assets.robots.diana_med import BiDianaMed
|
|
env = DualDianaMed_Pos_Ctrl(
|
|
robot=BiDianaMed(),
|
|
is_render=True,
|
|
control_freq=30,
|
|
is_interpolate=True,
|
|
cam_view='angle'
|
|
)
|
|
return env
|
|
else:
|
|
raise NotImplementedError
|
|
|
|
|
|
|
|
|
|
|
|
if __name__ == "__main__":
|
|
task_name = 'sim_transfer'
|
|
env = make_sim_env(task_name)
|
|
env.reset()
|
|
action_l = np.array([-0.20, 1.2, 1.1,
|
|
1.0, 0.0, 0.0, 0.0]) # 双臂直角
|
|
action_r = np.array([0.20, 1.2, 1.1,
|
|
1.0, 0.0, 0.0, 0.0])
|
|
action_gripper = np.array([-100,-100])
|
|
action = np.hstack((action_l,action_r,action_gripper))
|
|
|
|
for t in range(int(10000)):
|
|
# print("main step ", t)
|
|
env.step(action)
|
|
if env.is_render:
|
|
env.render()
|
|
|