chore(project): 项目初始化

roboimi/envs/double_pos_ctrl_env.py

@@ -0,0 +1,165 @@
+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.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)):
+                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()
+