<s> import numpy as np <EOL> from gym import utils <EOL> from gym . envs . mujoco import mujoco_env <EOL> class InvertedDoublePendulumEnv ( mujoco_env . MujocoEnv , utils . EzPickle ) : <EOL> def __init__ ( self ) : <EOL> mujoco_env . MujocoEnv . __init__ ( self , '<STR_LIT>' , <NUM_LIT:5> ) <EOL> utils . EzPickle . __init__ ( self ) <EOL> self . finalize ( ) <EOL> def _step ( self , action ) : <EOL> self . do_simulation ( action , self . frame_skip ) <EOL> ob = self . _get_obs ( ) <EOL> x , _ , y = self . model . data . site_xpos [ <NUM_LIT:0> ] <EOL> dist_penalty = <NUM_LIT> * x ** <NUM_LIT:2> + ( y - <NUM_LIT:2> ) ** <NUM_LIT:2> <EOL> v1 , v2 = self . model . data . qvel [ <NUM_LIT:1> : <NUM_LIT:3> ] <EOL> vel_penalty = <NUM_LIT> * v1 ** <NUM_LIT:2> + <NUM_LIT> * v2 ** <NUM_LIT:2> <EOL> alive_bonus = <NUM_LIT:10> <EOL> r = ( alive_bonus - dist_penalty - vel_penalty ) [ <NUM_LIT:0> ] <EOL> done = bool ( y <= <NUM_LIT:1> ) <EOL> return ob , r , done , { } <EOL> def _get_obs ( self ) : <EOL> return np . concatenate ( [ <EOL> self . model . data . qpos [ : <NUM_LIT:1> ] , <EOL> np . sin ( self . model . data . qpos [ <NUM_LIT:1> : ] ) , <EOL> np . cos ( self . model . data . qpos [ <NUM_LIT:1> : ] ) , <EOL> np . clip ( self . model . data . qvel , - <NUM_LIT:10> , <NUM_LIT:10> ) , <EOL> np . clip ( self . model . data . qfrc_constraint , - <NUM_LIT:10> , <NUM_LIT:10> ) <EOL> ] ) . ravel ( ) <EOL> def _reset ( self ) : <EOL> self . model . data . qpos = self . init_qpos + np . random . uniform ( size = ( self . model . nq , <NUM_LIT:1> ) , low = - <NUM_LIT> , high = <NUM_LIT> ) <EOL> self . model . data . qvel = self . init_qvel + np . random . randn ( self . model . nv , <NUM_LIT:1> ) * <NUM_LIT> <EOL> self . reset_viewer_if_necessary ( ) <EOL> return self . _get_obs ( ) <EOL> def viewer_setup ( self ) : <EOL> v = self . viewer <EOL> v . cam . trackbodyid = <NUM_LIT:0> <EOL> v . cam . distance = v . model . stat . extent * <NUM_LIT:0.5> <EOL> v . cam . lookat [ <NUM_LIT:2> ] += <NUM_LIT:3> </s>
