3 *]A@sddlmZddlZddlZddlZddlZddlmZddlm Z ddl Z ddl Z ddl Z ddl Z ddlZddlmZddlmZmZmZddlmZddlmZmZddlm Zdd lmZdd lmZdd lmZdd lm Z ddlm Z ddl!Z!d Z"d%Z#dZ$dZ%dZ&dZ'dZ(dZ)dZ*dZ+d&Z,d'Z-dZ.dZ/dZ0d(Z1dZ2e3e j4e)e,de*e-dZ5d)Z6d*Z7dZ8dZ9dZ:dZ;dZZ?Gd#d$d$e?Z@dS)+)print_functionN)spaces)Image)Odometry)TwistPose Quaternion) SetModelState) ModelState ContactsState) LaserScan)Point)Float64)Stringz0.0.4@g@gq= ףp?i'g!@g@g?g`l8~.?gaI?gڷlHb?ghKcR?g$@g?g?g333333?g{Gz?c@s|eZdZddZddZddZddZd d Zd d Zd dZ ddZ ddZ ddZ ddZ ddZddZddZdS)RoverTrainingGroundsEnvcCsht|_t|_d|_d|_d|_d|_tj |_ t|_ t|_ d|_ t|_tj |_d|_d|_d|_d|_d|_td|_tjtjddgtjddgtjd|_tjddtdtddftjd |_d|_t j!t"|_#t$j%d t&d d |_'t$j%d t(dd |_)t$j%dt*dd |_+t$j,dt-|_.t$j/dddt$j0dt1|j2t$j0dt3|j4t$j0dt5|j6t$j0dt7|j8dS)Nz us-east-1rFg@r)lowhighdtype)rrshaperz/cmd_veld) queue_sizez/current_positionz/distance_travelledz/gazebo/set_model_stateZrl_coachT)Z anonymousz/odomz/scanz /robot_bumperz/camera/image_rawrr)9 INITIAL_POS_Xx INITIAL_POS_Yy orientation aws_regionreward_in_episodestepssysmaxsizelast_collision_thresholdlast_position_xlast_position_ydistance_travelledINITIAL_DISTANCE_TO_CHECKPOINTcurrent_distance_to_checkpointcollision_threshold collisionstatesteeringthrottlecloser_to_checkpointpower_supply_rangerBoxnparrayfloat32 action_spaceTRAINING_IMAGE_SIZEuint8observation_spacequeueQueueIMG_QUEUE_BUF_SIZE image_queuerospyZ Publisherr ack_publisherr current_position_pubrZdistance_travelled_pubZ ServiceProxyr gazebo_model_state_serviceZ init_nodeZ Subscriberr callback_poser callback_scanr callback_collision sensor_imagecallback_image)selfrKW/home/ubuntu/environment/simulation_ws/src/rl-agent/markov/environments/training_env.py__init__MsB (  z RoverTrainingGroundsEnv.__init__cCsnd|_d|_d|_t|d}t|d}|jd7_|j||tjt|j |i}|j|j|j|fS)NFrr) rewarddone next_statefloatr% send_actiontimesleep:SLEEP_BETWEEN_ACTION_AND_REWARD_CALCULATION_TIME_IN_SECONDcall_reward_function)rJactionr1r2inforKrKrLsteps     zRoverTrainingGroundsEnv.stepcCsbtd|jd|j|j|jd|_d|_d|_d|_|jdd|j |j ddg|jS)NzTotal Episodic Reward=%.2fzTotal Episodic Steps=%.2fFr) printr$r%send_reward_to_cloudwatchrNrOrPrangesrR rover_resetrV)rJrKrKrLresets    zRoverTrainingGroundsEnv.resetcCs&t}||j_||j_|jj|dS)N)rlinearrangularzrBpublish)rJr1r2ZspeedrKrKrLrRsz#RoverTrainingGroundsEnv.send_actioncCstjdt|_t|_t}t|jj_t|jj_t |jj_ t |jj _t |jj _t|jj _ t|jj _d|jj_d|jj_d|jj_ d|jj_d|jj_d|jj_ d|_|j|tj|_|j|_|j|_tjtd|_t|_ d|_!d|_"d|_#d|_$td|_%|j&j'ddd}|j(dS)Nzgazebo/set_model_staterroverFg@T)blocktimeout))rAZwait_for_servicerrr r!r poseposition INITIAL_POS_ZraINITIAL_ORIENT_Xr"INITIAL_ORIENT_YINITIAL_ORIENT_ZINITIAL_ORIENT_Wwtwistr_r` model_namerDr&r'r(r)r*rSrT SLEEP_AFTER_RESET_TIME_IN_SECONDr+r,r-r%r$r/r3r4r@getset_next_state)rJZ model_state_rKrKrLr]s@                 z#RoverTrainingGroundsEnv.rover_resetcCstx|jstjtqWt|j}tjd|dd}tj|}tjtj |||jdt }t |tj |<tj ||_ dS)Nrrih)r\rSrT&SLEEP_WAITING_FOR_IMAGE_TIME_IN_SECONDlenr6linspacearangeclipinterpLIDAR_SCAN_MAX_DISTANCEisnanaminr.)rJsizerxpvalrKrKrLget_distance_to_objects  z.RoverTrainingGroundsEnv.get_distance_to_objectcCsyN|jjddd}tjd|j|jf|jdddd}|jtdd}t j ||_ Wn0t k r~}zt d j|WYdd}~XnXdS) NT)rdreRGBrawrrr)resamplezError!::set_next_state:: {})r@rqr frombyteswidthheightdataresizer:r6r7rP ExceptionrZformat)rJZ image_dataimageerrrKrKrLrrs z&RoverTrainingGroundsEnv.set_next_statec Cs|j|j|j\}}|j|7_td|jd|dd|d|jd|jd|jd|j d |j ||_ ||_ |j |_|j|_dS) Nz Step:%.2fz Steering:%frzR:%.2fzDTCP:%fzDT:%fzCT:%.2fzCTCP:%fzPSR: %f)rrrreward_functionr$rZr%r-r+r.r3r4rNrOrr)r!r*)rJrWrNrOrKrKrLrVs     z,RoverTrainingGroundsEnv.call_reward_functioncCsd}|jtkrD|jtkrDtd|jdkr4d}d}nt|j}d}n|jr\tdt}d}n|j t krxtdt}d}n|j t dks|j t dkrtd t}d}n`|jtdks|jtdkrtd t}d}n2|jd krtd t}d}n|jrd}nd}d}||fS)zK :return: reward as float done as boolean rz6Congratulations! The rover has reached the checkpoint!gFTz!Rover has collided with an objectz$Rover has sustained sideswipe damageg?zRover has left the mission map!rz%Rover's power supply has been drainedrg@g@gg)r) CHECKPOINT_Xr* CHECKPOINT_YrZr%FINISHEDr/COLLISION_REWARDr.CRASH_DISTANCEr STAGE_X_MIN STAGE_X_MAXr! STAGE_Y_MIN STAGE_Y_MAXr4r3)rJrNrOrKrKrLr<sB    z'RoverTrainingGroundsEnv.reward_functioncCs |j|_dS)N)r\)rJrrKrKrLrFnsz%RoverTrainingGroundsEnv.callback_scancCsZy|jj|WnDtjk r&Yn0tk rT}ztdj|WYdd}~XnXdS)Nz Error! {})r@ put_nowaitr=FullrrZr)rJrexrKrKrLrIrs z&RoverTrainingGroundsEnv.callback_imagecCs|jjj|_|jjj|_|jjj|_|jjj}t|j |j |j }|j j |tj|j |j |j |j }|j|7_|j|8_t}ttj|j td|j td|_|j|jkrd|_nd|_|j|_|j |_ |j |_ dS)NrTF)rfr"rnr_Zlinear_trajectoryr`Zangular_trajectoryrgr rr!rarCrbmathhypotr+r4rabssqrtrrrr-r3)rJrZ new_positionpdistZnew_distance_to_checkpointrKrKrLrE~s$      z%RoverTrainingGroundsEnv.callback_posecCs|j}t|dkrd|_dS)NrT)statesruZcollide)rJrZcollsion_statesrKrKrLrGs z*RoverTrainingGroundsEnv.callback_collisioncCsly6tjj}|jd|jd}|jdd|dgddWn0tk rf}ztdj|WYdd}~XnXdS) N cloudwatch) region_nameZTraining_GroundsNone)Z MetricNameZUnitValueZAWSRoboMakerSimulation)Z MetricData Namespacez3Error in the send_reward_to_cloudwatch function: {}) boto3sessionSessionclientr#Zput_metric_datarrZr)rJrNrZcloudwatch_clientrrKrKrLr[s  z1RoverTrainingGroundsEnv.send_reward_to_cloudwatchN)__name__ __module__ __qualname__rMrYr^rRr]rrrrVrrFrIrErGr[rKrKrKrLrLs6 -:2 ( rcs$eZdZddZfddZZS)RoverTrainingGroundsDiscreteEnvcCs"tj|tdtjd|_dS)Nz/New Training Grounds Gym environment created...r)rrMrZrDiscreter9)rJrKrKrLrMs z(RoverTrainingGroundsDiscreteEnv.__init__csR|dkrd}d}n,|dkr$d}d}n|dkr6d}d}ntd||g}tj|S)Nrg?g@rrzInvalid actiong) ValueErrorsuperrY)rJrWr1r2Zcontinuous_action) __class__rKrLrYsz$RoverTrainingGroundsDiscreteEnv.step)rrrrMrY __classcell__rKrK)rrLrsr)rrg!g!gڷlHbg$g$)A __future__rrSrgymnumpyr6rPILrosrandomrr&rAZ nav_msgs.msgrZgeometry_msgs.msgrrrZgazebo_msgs.srvr Zgazebo_msgs.msgr r Zsensor_msgs.msgrHr r Z std_msgs.msgrrr=VERSIONr:rzrr?r MAX_STEPSrrrrr rhrirjrkrlrrr,rrrrZFOOTSTEPS_MARKER_SIZErprUrtEnvrrrKrKrKrLsj           y