<s> import gym <EOL> from gym import spaces <EOL> import numpy as np <EOL> from os import path <EOL> class PendulumEnv ( gym . Env ) : <EOL> metadata = { <EOL> '<STR_LIT>' : [ '<STR_LIT>' , '<STR_LIT>' ] , <EOL> '<STR_LIT>' : <NUM_LIT:30> <EOL> } <EOL> def __init__ ( self ) : <EOL> self . max_speed = <NUM_LIT:8> <EOL> self . max_torque = <NUM_LIT> <EOL> self . dt = <NUM_LIT> <EOL> self . viewer = None <EOL> high = np . array ( [ <NUM_LIT:1.> , <NUM_LIT:1.> , self . max_speed ] ) <EOL> self . action_space = spaces . Box ( low = - self . max_torque , high = self . max_torque , shape = ( <NUM_LIT:1> , ) ) <EOL> self . observation_space = spaces . Box ( low = - high , high = high ) <EOL> def _step ( self , u ) : <EOL> th , thdot = self . state <EOL> g = <NUM_LIT> <EOL> m = <NUM_LIT:1.> <EOL> l = <NUM_LIT:1.> <EOL> dt = self . dt <EOL> self . last_u = u <EOL> u = np . clip ( u , - self . max_torque , self . max_torque ) [ <NUM_LIT:0> ] <EOL> costs = angle_normalize ( th ) ** <NUM_LIT:2> + <NUM_LIT> * thdot ** <NUM_LIT:2> + <NUM_LIT> * ( u ** <NUM_LIT:2> ) <EOL> newthdot = thdot + ( - <NUM_LIT:3> * g / ( <NUM_LIT:2> * l ) * np . sin ( th + np . pi ) + <NUM_LIT> / ( m * l ** <NUM_LIT:2> ) * u ) * dt <EOL> newth = th + newthdot * dt <EOL> newthdot = np . clip ( newthdot , - self . max_speed , self . max_speed ) <EOL> self . state = np . array ( [ newth , newthdot ] ) <EOL> return self . _get_obs ( ) , - costs , False , { } <EOL> def _reset ( self ) : <EOL> high = np . array ( [ np . pi , <NUM_LIT:1> ] ) <EOL> self . state = np . random . uniform ( low = - high , high = high ) <EOL> self . last_u = None <EOL> return self . _get_obs ( ) <EOL> def _get_obs ( self ) : <EOL> theta , thetadot = self . state <EOL> return np . array ( [ np . cos ( theta ) , np . sin ( theta ) , thetadot ] ) <EOL> def _render ( self , mode = '<STR_LIT>' , close = False ) : <EOL> if close : <EOL> if self . viewer is not None : <EOL> self . viewer . close ( ) <EOL> return <EOL> if self . viewer is None : <EOL> from gym . envs . classic_control import rendering <EOL> self . viewer = rendering . Viewer ( <NUM_LIT> , <NUM_LIT> ) <EOL> self . viewer . set_bounds ( - <NUM_LIT> , <NUM_LIT> , - <NUM_LIT> , <NUM_LIT> ) <EOL> rod = rendering . make_capsule ( <NUM_LIT:1> , <NUM_LIT> ) <EOL> rod . set_color ( <NUM_LIT> , <NUM_LIT> , <NUM_LIT> ) <EOL> self . pole_transform = rendering . Transform ( ) <EOL> rod . add_attr ( self . pole_transform ) <EOL> self . viewer . add_geom ( rod ) <EOL> axle = rendering . make_circle ( <NUM_LIT> ) <EOL> axle . set_color ( <NUM_LIT:0> , <NUM_LIT:0> , <NUM_LIT:0> ) <EOL> self . viewer . add_geom ( axle ) <EOL> fname = path . join ( path . dirname ( __file__ ) , "<STR_LIT>" ) <EOL> self . img = rendering . Image ( fname , <NUM_LIT:1.> , <NUM_LIT:1.> ) <EOL> self . imgtrans = rendering . Transform ( ) <EOL> self . img . add_attr ( self . imgtrans ) <EOL> self . viewer . add_onetime ( self . img ) <EOL> self . pole_transform . set_rotation ( self . state [ <NUM_LIT:0> ] + np . pi / <NUM_LIT:2> ) <EOL> if self . last_u : <EOL> self . imgtrans . scale = ( - self . last_u / <NUM_LIT:2> , np . abs ( self . last_u ) / <NUM_LIT:2> ) <EOL> self . viewer . render ( ) <EOL> if mode == '<STR_LIT>' : <EOL> return self . viewer . get_array ( ) <EOL> elif mode is '<STR_LIT>' : <EOL> pass <EOL> else : <EOL> return super ( PendulumEnv , self ) . render ( mode = mode ) <EOL> def angle_normalize ( x ) : <EOL> return ( ( ( x + np . pi ) % ( <NUM_LIT:2> * np . pi ) ) - np . pi ) </s>
