<s> import bulletsimpy <EOL> import numpy as np <EOL> from lfd . rapprentice import math_utils , retiming <EOL> import trajoptpy <EOL> def transform ( hmat , p ) : <EOL> return hmat [ : <NUM_LIT:3> , : <NUM_LIT:3> ] . dot ( p ) + hmat [ : <NUM_LIT:3> , <NUM_LIT:3> ] <EOL> def in_grasp_region ( robot , lr , pt ) : <EOL> tol = <NUM_LIT> <EOL> manip_name = { "<STR_LIT:l>" : "<STR_LIT>" , "<STR_LIT:r>" : "<STR_LIT>" } [ lr ] <EOL> manip = robot . GetManipulator ( manip_name ) <EOL> l_finger = robot . GetLink ( "<STR_LIT>" % lr ) <EOL> r_finger = robot . GetLink ( "<STR_LIT>" % lr ) <EOL> def on_inner_side ( pt , finger_lr ) : <EOL> finger = l_finger <EOL> closing_dir = np . cross ( manip . GetLocalToolDirection ( ) , [ - <NUM_LIT:1> , <NUM_LIT:0> , <NUM_LIT:0> ] ) <EOL> local_inner_pt = np . array ( [ <NUM_LIT> , - <NUM_LIT> , <NUM_LIT:0> ] ) / <NUM_LIT> <EOL> if finger_lr == "<STR_LIT:r>" : <EOL> finger = r_finger <EOL> closing_dir *= - <NUM_LIT:1> <EOL> local_inner_pt [ <NUM_LIT:1> ] *= - <NUM_LIT:1> <EOL> inner_pt = transform ( finger . GetTransform ( ) , local_inner_pt ) <EOL> return manip . GetTransform ( ) [ : <NUM_LIT:3> , : <NUM_LIT:3> ] . dot ( closing_dir ) . dot ( pt - inner_pt ) > <NUM_LIT:0> <EOL> pt_local = transform ( np . linalg . inv ( manip . GetTransform ( ) ) , pt ) <EOL> if pt_local [ <NUM_LIT:2> ] > <NUM_LIT> + tol : <EOL> return False <EOL> if abs ( pt_local [ <NUM_LIT:0> ] ) > <NUM_LIT> + tol : <EOL> return False <EOL> if not on_inner_side ( pt , "<STR_LIT:l>" ) or not on_inner_side ( pt , "<STR_LIT:r>" ) : <EOL> return False <EOL> return True <EOL> def retime_traj ( robot , inds , traj , max_cart_vel = <NUM_LIT> , max_finger_vel = <NUM_LIT> , upsample_time = <NUM_LIT> ) : <EOL> """<STR_LIT>""" <EOL> cart_traj = np . empty ( ( len ( traj ) , <NUM_LIT:6> ) ) <EOL> finger_traj = np . empty ( ( len ( traj ) , <NUM_LIT:2> ) ) <EOL> leftarm , rightarm = robot . GetManipulator ( "<STR_LIT>" ) , robot . GetManipulator ( "<STR_LIT>" ) <EOL> with robot : <EOL> for i in range ( len ( traj ) ) : <EOL> robot . SetDOFValues ( traj [ i ] , inds ) <EOL> cart_traj [ i , : <NUM_LIT:3> ] = leftarm . GetTransform ( ) [ : <NUM_LIT:3> , <NUM_LIT:3> ] <EOL> cart_traj [ i , <NUM_LIT:3> : ] = rightarm . GetTransform ( ) [ : <NUM_LIT:3> , <NUM_LIT:3> ] <EOL> finger_traj [ i , : <NUM_LIT:1> ] = robot . GetDOFValues ( leftarm . GetGripperIndices ( ) ) <EOL> finger_traj [ i , <NUM_LIT:1> : ] = robot . GetDOFValues ( rightarm . GetGripperIndices ( ) ) <EOL> times = retiming . retime_with_vel_limits ( np . c_ [ cart_traj , finger_traj ] , np . r_ [ np . repeat ( max_cart_vel , <NUM_LIT:6> ) , np . repeat ( max_finger_vel , <NUM_LIT:2> ) ] ) <EOL> times_up = np . linspace ( <NUM_LIT:0> , times [ - <NUM_LIT:1> ] , times [ - <NUM_LIT:1> ] / upsample_time ) if times [ - <NUM_LIT:1> ] > upsample_time else times <EOL> traj_up = math_utils . interp2d ( times_up , times , traj ) <EOL> return traj_up <EOL> def observe_cloud ( pts , radius , upsample = <NUM_LIT:0> , upsample_rad = <NUM_LIT:1> ) : <EOL> """<STR_LIT>""" <EOL> if upsample > <NUM_LIT:0> : <EOL> lengths = np . r_ [ <NUM_LIT:0> , np . apply_along_axis ( np . linalg . norm , <NUM_LIT:1> , np . diff ( pts , axis = <NUM_LIT:0> ) ) ] <EOL> summed_lengths = np . cumsum ( lengths ) <EOL> assert len ( lengths ) == len ( pts ) <EOL> pts = math_utils . interp2d ( np . linspace ( <NUM_LIT:0> , summed_lengths [ - <NUM_LIT:1> ] , upsample ) , summed_lengths , pts ) <EOL> if upsample_rad > <NUM_LIT:1> : <EOL> vs = np . diff ( pts , axis = <NUM_LIT:0> ) <EOL> vs /= np . apply_along_axis ( np . linalg . norm , <NUM_LIT:1> , vs ) [ : , None ] <EOL> perp_vs = np . c_ [ - vs [ : , <NUM_LIT:1> ] , vs [ : , <NUM_LIT:0> ] , np . zeros ( vs . shape [ <NUM_LIT:0> ] ) ] <EOL> perp_vs /= np . apply_along_axis ( np . linalg . norm , <NUM_LIT:1> , perp_vs ) [ : , None ] <EOL> vs = np . r_ [ vs , vs [ - <NUM_LIT:1> , : ] [ None , : ] ] <EOL> perp_vs = np . r_ [ perp_vs , perp_vs [ - <NUM_LIT:1> , : ] [ None , : ] ] <EOL> perp_pts = [ ] <EOL> from openravepy import matrixFromAxisAngle <EOL> for theta in np . linspace ( <NUM_LIT:0> , <NUM_LIT:2> * np . pi , upsample_rad , endpoint = False ) : <EOL> for ( center , rot_axis , perp_v ) in zip ( pts , vs , perp_vs ) : <EOL> rot = matrixFromAxisAngle ( rot_axis , theta ) [ : <NUM_LIT:3> , : <NUM_LIT:3> ] <EOL> perp_pts . append ( center + rot . T . dot ( radius * perp_v ) ) <EOL> pts = np . array ( perp_pts ) <EOL> return pts <EOL> class Simulation ( object ) : <EOL> def __init__ ( self , env , robot , rope_params ) : <EOL> self . env = env <EOL> self . robot = robot <EOL> self . bt_env = None <EOL> self . bt_robot = None <EOL> self . rope = None <EOL> self . rope_pts = None <EOL> self . rope_params = rope_params <EOL> self . constraints = { "<STR_LIT:l>" : [ ] , "<STR_LIT:r>" : [ ] } <EOL> self . constraints_links = { "<STR_LIT:l>" : [ ] , "<STR_LIT:r>" : [ ] } <EOL> def create ( self , rope_pts ) : <EOL> self . bt_env = bulletsimpy . BulletEnvironment ( self . env , [ ] ) <EOL> self . bt_env . SetGravity ( [ <NUM_LIT:0> , <NUM_LIT:0> , - <NUM_LIT> ] ) <EOL> self . bt_robot = self . bt_env . GetObjectByName ( self . robot . GetName ( ) ) <EOL> capsule_rope_params = bulletsimpy . CapsuleRopeParams ( ) <EOL> capsule_rope_params . radius = self . rope_params . radius <EOL> capsule_rope_params . angStiffness = self . rope_params . angStiffness <EOL> capsule_rope_params . angDamping = self . rope_params . angDamping <EOL> capsule_rope_params . linDamping = self . rope_params . linDamping <EOL> capsule_rope_params . angLimit = self . rope_params . angLimit <EOL> capsule_rope_params . linStopErp = self . rope_params . linStopErp <EOL> capsule_rope_params . mass = self . rope_params . mass <EOL> self . rope = bulletsimpy . CapsuleRope ( self . bt_env , '<STR_LIT>' , rope_pts , capsule_rope_params ) <EOL> self . rope_pts = rope_pts <EOL> cc = trajoptpy . GetCollisionChecker ( self . env ) <EOL> for gripper_link in [ link for link in self . robot . GetLinks ( ) if '<STR_LIT>' in link . GetName ( ) ] : <EOL> for rope_link in self . rope . GetKinBody ( ) . GetLinks ( ) : <EOL> cc . ExcludeCollisionPair ( gripper_link , rope_link ) <EOL> self . settle ( ) <EOL> def __del__ ( self ) : <EOL> if self . rope : <EOL> cc = trajoptpy . GetCollisionChecker ( self . env ) <EOL> for gripper_link in [ link for link in self . robot . GetLinks ( ) if '<STR_LIT>' in link . GetName ( ) ] : <EOL> for rope_link in self . rope . GetKinBody ( ) . GetLinks ( ) : <EOL> cc . IncludeCollisionPair ( gripper_link , rope_link ) <EOL> for rope_link0 in self . rope . GetKinBody ( ) . GetLinks ( ) : <EOL> for rope_link1 in self . rope . GetKinBody ( ) . GetLinks ( ) : <EOL> cc . IncludeCollisionPair ( rope_link0 , rope_link1 ) <EOL> self . env . Remove ( self . rope . GetKinBody ( ) ) <EOL> self . rope = None <EOL> self . rope_pts = None <EOL> def step ( self ) : <EOL> self . bt_robot . UpdateBullet ( ) <EOL> self . bt_env . Step ( <NUM_LIT> , <NUM_LIT:200> , <NUM_LIT> ) <EOL> self . rope . UpdateRave ( ) <EOL> self . env . UpdatePublishedBodies ( ) <EOL> def settle ( self , max_steps = <NUM_LIT:100> , tol = <NUM_LIT> , animate = False ) : <EOL> """<STR_LIT>""" <EOL> prev_nodes = self . rope . GetNodes ( ) <EOL> for i in range ( max_steps ) : <EOL> self . bt_env . Step ( <NUM_LIT> , <NUM_LIT:200> , <NUM_LIT> ) <EOL> if animate : <EOL> self . rope . UpdateRave ( ) <EOL> self . env . UpdatePublishedBodies ( ) <EOL> if i % <NUM_LIT:10> == <NUM_LIT:0> and i != <NUM_LIT:0> : <EOL> curr_nodes = self . rope . GetNodes ( ) <EOL> diff = np . sqrt ( ( ( curr_nodes - prev_nodes ) ** <NUM_LIT:2> ) . sum ( axis = <NUM_LIT:1> ) ) <EOL> if diff . max ( ) < tol : <EOL> break <EOL> prev_nodes = curr_nodes <EOL> self . rope . UpdateRave ( ) <EOL> self . env . UpdatePublishedBodies ( ) <EOL> def observe_cloud ( self , upsample = <NUM_LIT:0> , upsample_rad = <NUM_LIT:1> ) : <EOL> """<STR_LIT>""" <EOL> return observe_cloud ( self . rope . GetControlPoints ( ) , self . rope_params . radius , upsample = upsample , upsample_rad = upsample_rad ) <EOL> def raycast_cloud ( self , T_w_k = None , obj = None , z = <NUM_LIT:1.> , endpoints = <NUM_LIT:0> ) : <EOL> """<STR_LIT>""" <EOL> if T_w_k is None : <EOL> T_w_k = self . robot . GetLink ( "<STR_LIT>" ) . GetTransform ( ) <EOL> if obj is None : <EOL> obj = self . rope <EOL> cx = <NUM_LIT> - <NUM_LIT> <EOL> cy = <NUM_LIT> - <NUM_LIT> <EOL> f = <NUM_LIT> <EOL> w = <NUM_LIT> <EOL> h = <NUM_LIT> <EOL> pixel_ij = np . array ( np . meshgrid ( np . arange ( w ) , np . arange ( h ) ) ) . T . reshape ( ( - <NUM_LIT:1> , <NUM_LIT:2> ) ) <EOL> rayTos = z * np . c_ [ ( pixel_ij - np . array ( [ cx , cy ] ) ) / f , np . ones ( pixel_ij . shape [ <NUM_LIT:0> ] ) ] <EOL> rayFroms = np . zeros_like ( rayTos ) <EOL> rayTos = rayTos . dot ( T_w_k [ : <NUM_LIT:3> , : <NUM_LIT:3> ] . T ) + T_w_k [ : <NUM_LIT:3> , <NUM_LIT:3> ] <EOL> rayFroms = rayFroms . dot ( T_w_k [ : <NUM_LIT:3> , : <NUM_LIT:3> ] . T ) + T_w_k [ : <NUM_LIT:3> , <NUM_LIT:3> ] <EOL> ray_collisions = self . bt_env . RayTest ( rayFroms , rayTos , obj ) <EOL> pts = np . empty ( ( len ( ray_collisions ) , <NUM_LIT:3> ) ) <EOL> for i , ray_collision in enumerate ( ray_collisions ) : <EOL> pts [ i , : ] = ray_collision . pt <EOL> if endpoints : <EOL> links = self . rope . GetKinBody ( ) . GetLinks ( ) <EOL> end_links = links [ : endpoints ] + links [ - endpoints : ] <EOL> endpoint_inds = np . zeros ( len ( pts ) , dtype = bool ) <EOL> for i , ray_collision in enumerate ( ray_collisions ) : <EOL> if ray_collision . link in end_links : <EOL> endpoint_inds [ i ] = True <EOL> return pts , endpoint_inds <EOL> return pts <EOL> def grab_rope ( self , lr ) : <EOL> nodes , ctl_pts = self . rope . GetNodes ( ) , self . rope . GetControlPoints ( ) <EOL> graspable_nodes = np . array ( [ in_grasp_region ( self . robot , lr , n ) for n in nodes ] ) <EOL> graspable_ctl_pts = np . array ( [ in_grasp_region ( self . robot , lr , n ) for n in ctl_pts ] ) <EOL> graspable_inds = np . flatnonzero ( np . logical_or ( graspable_nodes , np . logical_or ( graspable_ctl_pts [ : - <NUM_LIT:1> ] , graspable_ctl_pts [ <NUM_LIT:1> : ] ) ) ) <EOL> print '<STR_LIT>' % ( lr , str ( graspable_inds ) ) <EOL> if len ( graspable_inds ) == <NUM_LIT:0> : <EOL> return False <EOL> robot_link = self . robot . GetLink ( "<STR_LIT>" % lr ) <EOL> rope_links = self . rope . GetKinBody ( ) . GetLinks ( ) <EOL> for i_node in graspable_inds : <EOL> i_cnt = i_node <EOL> for geom in rope_links [ i_cnt ] . GetGeometries ( ) : <EOL> geom . SetDiffuseColor ( [ <NUM_LIT:1.> , <NUM_LIT:0.> , <NUM_LIT:0.> ] ) <EOL> cnt = self . bt_env . AddConstraint ( { <EOL> "<STR_LIT:type>" : "<STR_LIT>" , <EOL> "<STR_LIT>" : { <EOL> "<STR_LIT>" : robot_link , <EOL> "<STR_LIT>" : rope_links [ i_cnt ] , <EOL> "<STR_LIT>" : np . linalg . inv ( robot_link . GetTransform ( ) ) . dot ( rope_links [ i_cnt ] . GetTransform ( ) ) , <EOL> "<STR_LIT>" : np . eye ( <NUM_LIT:4> ) , <EOL> "<STR_LIT>" : False , <EOL> "<STR_LIT>" : <NUM_LIT> , <EOL> "<STR_LIT>" : <NUM_LIT> , <EOL> "<STR_LIT>" : True , <EOL> } <EOL> } ) <EOL> self . constraints [ lr ] . append ( cnt ) <EOL> self . constraints_links [ lr ] . append ( rope_links [ i_cnt ] ) <EOL> return True <EOL> def release_rope ( self , lr ) : <EOL> for c in self . constraints [ lr ] : <EOL> self . bt_env . RemoveConstraint ( c ) <EOL> rope_links = self . rope . GetKinBody ( ) . GetLinks ( ) <EOL> for link in self . constraints_links [ lr ] : <EOL> for geom in link . GetGeometries ( ) : <EOL> geom . SetDiffuseColor ( [ <NUM_LIT:1.> , <NUM_LIT:1.> , <NUM_LIT:1.> ] ) <EOL> self . constraints [ lr ] = [ ] <EOL> self . constraints_links [ lr ] = [ ] <EOL> def is_grabbing_rope ( self , lr ) : <EOL> return bool ( self . constraints [ lr ] ) </s>
