<s> import os . path as osp <EOL> import mako . lookup <EOL> import mako . template <EOL> import numpy as np <EOL> from rllab import spaces <EOL> from rllab . envs . base import Env , Step <EOL> from rllab . envs . box2d . box2d_viewer import Box2DViewer <EOL> from rllab . envs . box2d . parser . xml_box2d import world_from_xml , find_body , find_joint <EOL> from rllab . misc import autoargs <EOL> from rllab . misc . overrides import overrides <EOL> BIG = <NUM_LIT> <EOL> class Box2DEnv ( Env ) : <EOL> @ autoargs . arg ( "<STR_LIT>" , type = int , <EOL> help = "<STR_LIT>" ) <EOL> @ autoargs . arg ( '<STR_LIT>' , type = bool , <EOL> help = '<STR_LIT>' <EOL> '<STR_LIT>' ) <EOL> @ autoargs . arg ( '<STR_LIT>' , type = float , <EOL> help = '<STR_LIT>' <EOL> '<STR_LIT>' ) <EOL> @ autoargs . arg ( '<STR_LIT>' , type = float , <EOL> help = '<STR_LIT>' <EOL> '<STR_LIT>' ) <EOL> def __init__ ( <EOL> self , model_path , frame_skip = <NUM_LIT:1> , position_only = False , <EOL> obs_noise = <NUM_LIT:0.0> , action_noise = <NUM_LIT:0.0> , template_string = None , <EOL> template_args = None , <EOL> ) : <EOL> self . full_model_path = model_path <EOL> if template_string is None : <EOL> if model_path . endswith ( "<STR_LIT>" ) : <EOL> with open ( model_path ) as template_file : <EOL> template = mako . template . Template ( <EOL> template_file . read ( ) ) <EOL> template_string = template . render ( <EOL> opts = template_args if template_args is not None else { } , <EOL> ) <EOL> else : <EOL> with open ( model_path , "<STR_LIT:r>" ) as f : <EOL> template_string = f . read ( ) <EOL> world , extra_data = world_from_xml ( template_string ) <EOL> self . world = world <EOL> self . extra_data = extra_data <EOL> self . initial_state = self . _state <EOL> self . viewer = None <EOL> self . frame_skip = frame_skip <EOL> self . timestep = self . extra_data . timeStep <EOL> self . position_only = position_only <EOL> self . obs_noise = obs_noise <EOL> self . action_noise = action_noise <EOL> self . _action_bounds = None <EOL> self . _position_ids = None <EOL> self . _cached_obs = None <EOL> self . _cached_coms = { } <EOL> def model_path ( self , file_name ) : <EOL> return osp . abspath ( osp . join ( osp . dirname ( __file__ ) , <EOL> '<STR_LIT>' % file_name ) ) <EOL> def _set_state ( self , state ) : <EOL> splitted = np . array ( state ) . reshape ( ( - <NUM_LIT:1> , <NUM_LIT:6> ) ) <EOL> for body , body_state in zip ( self . world . bodies , splitted ) : <EOL> xpos , ypos , apos , xvel , yvel , avel = body_state <EOL> body . position = ( xpos , ypos ) <EOL> body . angle = apos <EOL> body . linearVelocity = ( xvel , yvel ) <EOL> body . angularVelocity = avel <EOL> @ overrides <EOL> def reset ( self ) : <EOL> self . _set_state ( self . initial_state ) <EOL> self . _invalidate_state_caches ( ) <EOL> return self . get_current_obs ( ) <EOL> def _invalidate_state_caches ( self ) : <EOL> self . _cached_obs = None <EOL> self . _cached_coms = { } <EOL> @ property <EOL> def _state ( self ) : <EOL> s = [ ] <EOL> for body in self . world . bodies : <EOL> s . append ( np . concatenate ( [ <EOL> list ( body . position ) , <EOL> [ body . angle ] , <EOL> list ( body . linearVelocity ) , <EOL> [ body . angularVelocity ] <EOL> ] ) ) <EOL> return np . concatenate ( s ) <EOL> @ property <EOL> @ overrides <EOL> def action_space ( self ) : <EOL> lb = np . array ( [ control . ctrllimit [ <NUM_LIT:0> ] for control in self . extra_data . controls ] ) <EOL> ub = np . array ( [ control . ctrllimit [ <NUM_LIT:1> ] for control in self . extra_data . controls ] ) <EOL> return spaces . Box ( lb , ub ) <EOL> @ property <EOL> @ overrides <EOL> def observation_space ( self ) : <EOL> if self . position_only : <EOL> d = len ( self . _get_position_ids ( ) ) <EOL> else : <EOL> d = len ( self . extra_data . states ) <EOL> ub = BIG * np . ones ( d ) <EOL> return spaces . Box ( ub * - <NUM_LIT:1> , ub ) <EOL> @ property <EOL> def action_bounds ( self ) : <EOL> return self . action_space . bounds <EOL> def forward_dynamics ( self , action ) : <EOL> if len ( action ) != self . action_dim : <EOL> raise ValueError ( '<STR_LIT>' <EOL> '<STR_LIT>' % ( self . action_dim , len ( action ) ) ) <EOL> lb , ub = self . action_bounds <EOL> action = np . clip ( action , lb , ub ) <EOL> for ctrl , act in zip ( self . extra_data . controls , action ) : <EOL> if ctrl . typ == "<STR_LIT>" : <EOL> for name in ctrl . bodies : <EOL> body = find_body ( self . world , name ) <EOL> direction = np . array ( ctrl . direction ) <EOL> direction = direction / np . linalg . norm ( direction ) <EOL> world_force = body . GetWorldVector ( direction * act ) <EOL> world_point = body . GetWorldPoint ( ctrl . anchor ) <EOL> body . ApplyForce ( world_force , world_point , wake = True ) <EOL> elif ctrl . typ == "<STR_LIT>" : <EOL> assert ctrl . joint <EOL> joint = find_joint ( self . world , ctrl . joint ) <EOL> joint . motorEnabled = True <EOL> if act > <NUM_LIT:0> : <EOL> joint . motorSpeed = <NUM_LIT> <EOL> else : <EOL> joint . motorSpeed = - <NUM_LIT> <EOL> joint . maxMotorTorque = abs ( act ) <EOL> else : <EOL> raise NotImplementedError <EOL> self . before_world_step ( action ) <EOL> self . world . Step ( <EOL> self . extra_data . timeStep , <EOL> self . extra_data . velocityIterations , <EOL> self . extra_data . positionIterations <EOL> ) <EOL> def compute_reward ( self , action ) : <EOL> """<STR_LIT>""" <EOL> raise NotImplementedError <EOL> @ overrides <EOL> def step ( self , action ) : <EOL> """<STR_LIT>""" <EOL> reward_computer = self . compute_reward ( action ) <EOL> action = self . _inject_action_noise ( action ) <EOL> for _ in range ( self . frame_skip ) : <EOL> self . forward_dynamics ( action ) <EOL> reward_computer . next ( ) <EOL> reward = reward_computer . next ( ) <EOL> self . _invalidate_state_caches ( ) <EOL> done = self . is_current_done ( ) <EOL> next_obs = self . get_current_obs ( ) <EOL> return Step ( observation = next_obs , reward = reward , done = done ) <EOL> def _filter_position ( self , obs ) : <EOL> """<STR_LIT>""" <EOL> return obs [ self . _get_position_ids ( ) ] <EOL> def get_obs_noise_scale_factor ( self , obs ) : <EOL> return np . ones_like ( obs ) <EOL> def _inject_obs_noise ( self , obs ) : <EOL> """<STR_LIT>""" <EOL> noise = self . get_obs_noise_scale_factor ( obs ) * self . obs_noise * np . random . normal ( size = obs . shape ) <EOL> return obs + noise <EOL> def get_current_reward ( <EOL> self , state , xml_obs , action , next_state , next_xml_obs ) : <EOL> raise NotImplementedError <EOL> def is_current_done ( self ) : <EOL> raise NotImplementedError <EOL> def _inject_action_noise ( self , action ) : <EOL> noise = self . action_noise * np . random . normal ( size = action . shape ) <EOL> lb , ub = self . action_bounds <EOL> noise = <NUM_LIT:0.5> * ( ub - lb ) * noise <EOL> return action + noise <EOL> def get_current_obs ( self ) : <EOL> """<STR_LIT>""" <EOL> raw_obs = self . get_raw_obs ( ) <EOL> noisy_obs = self . _inject_obs_noise ( raw_obs ) <EOL> if self . position_only : <EOL> return self . _filter_position ( noisy_obs ) <EOL> return noisy_obs <EOL> def _get_position_ids ( self ) : <EOL> if self . _position_ids is None : <EOL> self . _position_ids = [ ] <EOL> for idx , state in enumerate ( self . extra_data . states ) : <EOL> if state . typ in [ "<STR_LIT>" , "<STR_LIT>" , "<STR_LIT>" , "<STR_LIT>" , "<STR_LIT>" ] : <EOL> self . _position_ids . append ( idx ) <EOL> return self . _position_ids <EOL> def get_raw_obs ( self ) : <EOL> """<STR_LIT>""" <EOL> if self . _cached_obs is not None : <EOL> return self . _cached_obs <EOL> obs = [ ] <EOL> for state in self . extra_data . states : <EOL> new_obs = None <EOL> if state . body : <EOL> body = find_body ( self . world , state . body ) <EOL> if state . local is not None : <EOL> l = state . local <EOL> position = body . GetWorldPoint ( l ) <EOL> linearVel = body . GetLinearVelocityFromLocalPoint ( l ) <EOL> else : <EOL> position = body . position <EOL> linearVel = body . linearVelocity <EOL> if state . to is not None : <EOL> to = find_body ( self . world , state . to ) <EOL> if state . typ == "<STR_LIT>" : <EOL> new_obs = position [ <NUM_LIT:0> ] <EOL> elif state . typ == "<STR_LIT>" : <EOL> new_obs = position [ <NUM_LIT:1> ] <EOL> elif state . typ == "<STR_LIT>" : <EOL> new_obs = linearVel [ <NUM_LIT:0> ] <EOL> elif state . typ == "<STR_LIT>" : <EOL> new_obs = linearVel [ <NUM_LIT:1> ] <EOL> elif state . typ == "<STR_LIT>" : <EOL> new_obs = body . angle <EOL> elif state . typ == "<STR_LIT>" : <EOL> new_obs = body . angularVelocity <EOL> elif state . typ == "<STR_LIT>" : <EOL> new_obs = np . linalg . norm ( position - to . position ) <EOL> elif state . typ == "<STR_LIT>" : <EOL> diff = to . position - position <EOL> abs_angle = np . arccos ( <EOL> diff . dot ( ( <NUM_LIT:0> , <NUM_LIT:1> ) ) / np . linalg . norm ( diff ) ) <EOL> new_obs = body . angle + abs_angle <EOL> else : <EOL> raise NotImplementedError <EOL> elif state . joint : <EOL> joint = find_joint ( self . world , state . joint ) <EOL> if state . typ == "<STR_LIT>" : <EOL> new_obs = joint . angle <EOL> elif state . typ == "<STR_LIT>" : <EOL> new_obs = joint . speed <EOL> else : <EOL> raise NotImplementedError <EOL> elif state . com : <EOL> com_quant = self . _compute_com_pos_vel ( * state . com ) <EOL> if state . typ == "<STR_LIT>" : <EOL> new_obs = com_quant [ <NUM_LIT:0> ] <EOL> elif state . typ == "<STR_LIT>" : <EOL> new_obs = com_quant [ <NUM_LIT:1> ] <EOL> elif state . typ == "<STR_LIT>" : <EOL> new_obs = com_quant [ <NUM_LIT:2> ] <EOL> elif state . typ == "<STR_LIT>" : <EOL> new_obs = com_quant [ <NUM_LIT:3> ] <EOL> else : <EOL> print state . typ <EOL> raise NotImplementedError <EOL> else : <EOL> raise NotImplementedError <EOL> if state . transform is not None : <EOL> if state . transform == "<STR_LIT:id>" : <EOL> pass <EOL> elif state . transform == "<STR_LIT>" : <EOL> new_obs = np . sin ( new_obs ) <EOL> elif state . transform == "<STR_LIT>" : <EOL> new_obs = np . cos ( new_obs ) <EOL> else : <EOL> raise NotImplementedError <EOL> obs . append ( new_obs ) <EOL> self . _cached_obs = np . array ( obs ) <EOL> return self . _cached_obs <EOL> def _compute_com_pos_vel ( self , * com ) : <EOL> com_key = "<STR_LIT:U+002C>" . join ( sorted ( com ) ) <EOL> if com_key in self . _cached_coms : <EOL> return self . _cached_coms [ com_key ] <EOL> total_mass_quant = <NUM_LIT:0> <EOL> total_mass = <NUM_LIT:0> <EOL> for body_name in com : <EOL> body = find_body ( self . world , body_name ) <EOL> total_mass_quant += body . mass * np . array ( <EOL> list ( body . worldCenter ) + list ( body . linearVelocity ) ) <EOL> total_mass += body . mass <EOL> com_quant = total_mass_quant / total_mass <EOL> self . _cached_coms [ com_key ] = com_quant <EOL> return com_quant <EOL> def get_com_position ( self , * com ) : <EOL> return self . _compute_com_pos_vel ( * com ) [ : <NUM_LIT:2> ] <EOL> def get_com_velocity ( self , * com ) : <EOL> return self . _compute_com_pos_vel ( * com ) [ <NUM_LIT:2> : ] <EOL> @ overrides <EOL> def render ( self , states = None , actions = None , pause = False ) : <EOL> if not self . viewer : <EOL> self . viewer = Box2DViewer ( self . world ) <EOL> if states or actions or pause : <EOL> raise NotImplementedError <EOL> if not self . viewer : <EOL> self . start_viewer ( ) <EOL> if self . viewer : <EOL> self . viewer . loop_once ( ) <EOL> def before_world_step ( self , action ) : <EOL> pass <EOL> def action_from_keys ( self , keys ) : <EOL> raise NotImplementedError </s>
