<s> import os . path as osp <EOL> import tempfile <EOL> import xml . etree . ElementTree as ET <EOL> import math <EOL> import numpy as np <EOL> from rllab import spaces <EOL> from rllab . envs . base import Step <EOL> from rllab . envs . proxy_env import ProxyEnv <EOL> from rllab . core . serializable import Serializable <EOL> from rllab . envs . mujoco . mujoco_env import MODEL_DIR , BIG <EOL> from rllab . misc . overrides import overrides <EOL> from rllab . envs . mujoco . maze . maze_env_utils import ray_segment_intersect , point_distance <EOL> class MazeEnv ( ProxyEnv , Serializable ) : <EOL> MODEL_CLASS = None <EOL> ORI_IND = None <EOL> MAZE_HEIGHT = None <EOL> MAZE_SIZE_SCALING = None <EOL> MAZE_MAKE_CONTACTS = False <EOL> MAZE_STRUCTURE = [ <EOL> [ <NUM_LIT:1> , <NUM_LIT:1> , <NUM_LIT:1> , <NUM_LIT:1> , <NUM_LIT:1> ] , <EOL> [ <NUM_LIT:1> , '<STR_LIT:r>' , <NUM_LIT:0> , <NUM_LIT:0> , <NUM_LIT:1> ] , <EOL> [ <NUM_LIT:1> , <NUM_LIT:1> , <NUM_LIT:1> , <NUM_LIT:0> , <NUM_LIT:1> ] , <EOL> [ <NUM_LIT:1> , '<STR_LIT:g>' , <NUM_LIT:0> , <NUM_LIT:0> , <NUM_LIT:1> ] , <EOL> [ <NUM_LIT:1> , <NUM_LIT:1> , <NUM_LIT:1> , <NUM_LIT:1> , <NUM_LIT:1> ] , <EOL> ] <EOL> MANUAL_COLLISION = False <EOL> def __init__ ( <EOL> self , <EOL> n_bins = <NUM_LIT:20> , <EOL> sensor_range = <NUM_LIT> , <EOL> sensor_span = math . pi , <EOL> * args , <EOL> ** kwargs ) : <EOL> self . _n_bins = n_bins <EOL> self . _sensor_range = sensor_range <EOL> self . _sensor_span = sensor_span <EOL> model_cls = self . __class__ . MODEL_CLASS <EOL> if model_cls is None : <EOL> raise "<STR_LIT>" <EOL> xml_path = osp . join ( MODEL_DIR , model_cls . FILE ) <EOL> tree = ET . parse ( xml_path ) <EOL> worldbody = tree . find ( "<STR_LIT>" ) <EOL> size_scaling = self . __class__ . MAZE_SIZE_SCALING <EOL> height = self . __class__ . MAZE_HEIGHT <EOL> structure = self . __class__ . MAZE_STRUCTURE <EOL> torso_x , torso_y = self . _find_robot ( ) <EOL> self . _init_torso_x = torso_x <EOL> self . _init_torso_y = torso_y <EOL> for i in xrange ( len ( structure ) ) : <EOL> for j in xrange ( len ( structure [ <NUM_LIT:0> ] ) ) : <EOL> if str ( structure [ i ] [ j ] ) == '<STR_LIT:1>' : <EOL> ET . SubElement ( <EOL> worldbody , "<STR_LIT>" , <EOL> name = "<STR_LIT>" % ( i , j ) , <EOL> pos = "<STR_LIT>" % ( j * size_scaling - torso_x , <EOL> i * size_scaling - torso_y , <EOL> height / <NUM_LIT:2> * size_scaling ) , <EOL> size = "<STR_LIT>" % ( <NUM_LIT:0.5> * size_scaling , <EOL> <NUM_LIT:0.5> * size_scaling , <EOL> height / <NUM_LIT:2> * size_scaling ) , <EOL> type = "<STR_LIT>" , <EOL> material = "<STR_LIT>" , <EOL> contype = "<STR_LIT:1>" , <EOL> conaffinity = "<STR_LIT:1>" , <EOL> rgba = "<STR_LIT>" <EOL> ) <EOL> torso = tree . find ( "<STR_LIT>" ) <EOL> geoms = torso . findall ( "<STR_LIT>" ) <EOL> for geom in geoms : <EOL> if '<STR_LIT:name>' not in geom . attrib : <EOL> raise Exception ( "<STR_LIT>" <EOL> "<STR_LIT>" ) <EOL> if self . __class__ . MAZE_MAKE_CONTACTS : <EOL> contact = ET . SubElement ( <EOL> tree . find ( "<STR_LIT:.>" ) , "<STR_LIT>" <EOL> ) <EOL> for i in range ( len ( structure ) ) : <EOL> for j in range ( len ( structure [ <NUM_LIT:0> ] ) ) : <EOL> if str ( structure [ i ] [ j ] ) == '<STR_LIT:1>' : <EOL> for geom in geoms : <EOL> ET . SubElement ( <EOL> contact , "<STR_LIT>" , <EOL> geom1 = geom . attrib [ "<STR_LIT:name>" ] , <EOL> geom2 = "<STR_LIT>" % ( i , j ) <EOL> ) <EOL> _ , file_path = tempfile . mkstemp ( text = True ) <EOL> tree . write ( file_path ) <EOL> self . _goal_range = self . _find_goal_range ( ) <EOL> self . _cached_segments = None <EOL> inner_env = model_cls ( * args , file_path = file_path , ** kwargs ) <EOL> ProxyEnv . __init__ ( self , inner_env ) <EOL> Serializable . quick_init ( self , locals ( ) ) <EOL> def get_current_obs ( self ) : <EOL> robot_x , robot_y = self . wrapped_env . get_body_com ( "<STR_LIT>" ) [ : <NUM_LIT:2> ] <EOL> ori = self . wrapped_env . model . data . qpos [ self . __class__ . ORI_IND ] <EOL> structure = self . __class__ . MAZE_STRUCTURE <EOL> size_scaling = self . __class__ . MAZE_SIZE_SCALING <EOL> segments = [ ] <EOL> for i in range ( len ( structure ) ) : <EOL> for j in range ( len ( structure [ <NUM_LIT:0> ] ) ) : <EOL> if structure [ i ] [ j ] == <NUM_LIT:1> or structure [ i ] [ j ] == '<STR_LIT:g>' : <EOL> cx = j * size_scaling - self . _init_torso_x <EOL> cy = i * size_scaling - self . _init_torso_y <EOL> x1 = cx - <NUM_LIT:0.5> * size_scaling <EOL> x2 = cx + <NUM_LIT:0.5> * size_scaling <EOL> y1 = cy - <NUM_LIT:0.5> * size_scaling <EOL> y2 = cy + <NUM_LIT:0.5> * size_scaling <EOL> struct_segments = [ <EOL> ( ( x1 , y1 ) , ( x2 , y1 ) ) , <EOL> ( ( x2 , y1 ) , ( x2 , y2 ) ) , <EOL> ( ( x2 , y2 ) , ( x1 , y2 ) ) , <EOL> ( ( x1 , y2 ) , ( x1 , y1 ) ) , <EOL> ] <EOL> for seg in struct_segments : <EOL> segments . append ( dict ( <EOL> segment = seg , <EOL> type = structure [ i ] [ j ] , <EOL> ) ) <EOL> wall_readings = np . zeros ( self . _n_bins ) <EOL> goal_readings = np . zeros ( self . _n_bins ) <EOL> for ray_idx in xrange ( self . _n_bins ) : <EOL> ray_ori = ori - self . _sensor_span * <NUM_LIT:0.5> + <NUM_LIT:1.0> * ( <NUM_LIT:2> * ray_idx + <NUM_LIT:1> ) / ( <NUM_LIT:2> * self . _n_bins ) * self . _sensor_span <EOL> ray_segments = [ ] <EOL> for seg in segments : <EOL> p = ray_segment_intersect ( ray = ( ( robot_x , robot_y ) , ray_ori ) , segment = seg [ "<STR_LIT>" ] ) <EOL> if p is not None : <EOL> ray_segments . append ( dict ( <EOL> segment = seg [ "<STR_LIT>" ] , <EOL> type = seg [ "<STR_LIT:type>" ] , <EOL> ray_ori = ray_ori , <EOL> distance = point_distance ( p , ( robot_x , robot_y ) ) , <EOL> ) ) <EOL> if len ( ray_segments ) > <NUM_LIT:0> : <EOL> first_seg = sorted ( ray_segments , key = lambda x : x [ "<STR_LIT>" ] ) [ <NUM_LIT:0> ] <EOL> if first_seg [ "<STR_LIT:type>" ] == <NUM_LIT:1> : <EOL> if first_seg [ "<STR_LIT>" ] <= self . _sensor_range : <EOL> wall_readings [ ray_idx ] = ( self . _sensor_range - first_seg [ "<STR_LIT>" ] ) / self . _sensor_range <EOL> elif first_seg [ "<STR_LIT:type>" ] == '<STR_LIT:g>' : <EOL> if first_seg [ "<STR_LIT>" ] <= self . _sensor_range : <EOL> goal_readings [ ray_idx ] = ( self . _sensor_range - first_seg [ "<STR_LIT>" ] ) / self . _sensor_range <EOL> else : <EOL> assert False <EOL> obs = np . concatenate ( [ <EOL> self . wrapped_env . get_current_obs ( ) , <EOL> wall_readings , <EOL> goal_readings <EOL> ] ) <EOL> return obs <EOL> def reset ( self ) : <EOL> self . wrapped_env . reset ( ) <EOL> return self . get_current_obs ( ) <EOL> @ property <EOL> def viewer ( self ) : <EOL> return self . wrapped_env . viewer <EOL> @ property <EOL> @ overrides <EOL> def observation_space ( self ) : <EOL> shp = self . get_current_obs ( ) . shape <EOL> ub = BIG * np . ones ( shp ) <EOL> return spaces . Box ( ub * - <NUM_LIT:1> , ub ) <EOL> def _find_robot ( self ) : <EOL> structure = self . __class__ . MAZE_STRUCTURE <EOL> size_scaling = self . __class__ . MAZE_SIZE_SCALING <EOL> for i in range ( len ( structure ) ) : <EOL> for j in range ( len ( structure [ <NUM_LIT:0> ] ) ) : <EOL> if structure [ i ] [ j ] == '<STR_LIT:r>' : <EOL> return j * size_scaling , i * size_scaling <EOL> assert False <EOL> def _find_goal_range ( self ) : <EOL> structure = self . __class__ . MAZE_STRUCTURE <EOL> size_scaling = self . __class__ . MAZE_SIZE_SCALING <EOL> for i in xrange ( len ( structure ) ) : <EOL> for j in xrange ( len ( structure [ <NUM_LIT:0> ] ) ) : <EOL> if structure [ i ] [ j ] == '<STR_LIT:g>' : <EOL> minx = j * size_scaling - size_scaling * <NUM_LIT:0.5> - self . _init_torso_x <EOL> maxx = j * size_scaling + size_scaling * <NUM_LIT:0.5> - self . _init_torso_x <EOL> miny = i * size_scaling - size_scaling * <NUM_LIT:0.5> - self . _init_torso_y <EOL> maxy = i * size_scaling + size_scaling * <NUM_LIT:0.5> - self . _init_torso_y <EOL> return minx , maxx , miny , maxy <EOL> def _is_in_collision ( self , pos ) : <EOL> x , y = pos <EOL> structure = self . __class__ . MAZE_STRUCTURE <EOL> size_scaling = self . __class__ . MAZE_SIZE_SCALING <EOL> for i in xrange ( len ( structure ) ) : <EOL> for j in xrange ( len ( structure [ <NUM_LIT:0> ] ) ) : <EOL> if structure [ i ] [ j ] == <NUM_LIT:1> : <EOL> minx = j * size_scaling - size_scaling * <NUM_LIT:0.5> - self . _init_torso_x <EOL> maxx = j * size_scaling + size_scaling * <NUM_LIT:0.5> - self . _init_torso_x <EOL> miny = i * size_scaling - size_scaling * <NUM_LIT:0.5> - self . _init_torso_y <EOL> maxy = i * size_scaling + size_scaling * <NUM_LIT:0.5> - self . _init_torso_y <EOL> if minx <= x <= maxx and miny <= y <= maxy : <EOL> return True <EOL> return False <EOL> def step ( self , action ) : <EOL> if self . MANUAL_COLLISION : <EOL> old_pos = self . wrapped_env . get_xy ( ) <EOL> _ , _ , done , info = self . wrapped_env . step ( action ) <EOL> new_pos = self . wrapped_env . get_xy ( ) <EOL> if self . _is_in_collision ( new_pos ) : <EOL> self . wrapped_env . set_xy ( old_pos ) <EOL> done = False <EOL> else : <EOL> _ , _ , done , info = self . wrapped_env . step ( action ) <EOL> next_obs = self . get_current_obs ( ) <EOL> x , y = self . wrapped_env . get_body_com ( "<STR_LIT>" ) [ : <NUM_LIT:2> ] <EOL> reward = <NUM_LIT:0> <EOL> minx , maxx , miny , maxy = self . _goal_range <EOL> if minx <= x <= maxx and miny <= y <= maxy : <EOL> done = True <EOL> reward = <NUM_LIT:1> <EOL> return Step ( next_obs , reward , done , ** info ) <EOL> def action_from_key ( self , key ) : <EOL> return self . wrapped_env . action_from_key ( key ) </s>
