<s> import numpy as np <EOL> import os . path as osp <EOL> from rllab import spaces <EOL> from rllab . envs . base import Env <EOL> from rllab . misc . overrides import overrides <EOL> from rllab . mujoco_py import MjModel , MjViewer <EOL> from rllab . misc import autoargs <EOL> from rllab . misc import logger <EOL> import theano <EOL> import tempfile <EOL> import os <EOL> import mako . template <EOL> import mako . lookup <EOL> MODEL_DIR = osp . abspath ( <EOL> osp . join ( <EOL> osp . dirname ( __file__ ) , <EOL> '<STR_LIT>' <EOL> ) <EOL> ) <EOL> BIG = <NUM_LIT> <EOL> class MujocoEnv ( Env ) : <EOL> FILE = None <EOL> @ autoargs . arg ( '<STR_LIT>' , type = float , <EOL> help = '<STR_LIT>' <EOL> '<STR_LIT>' ) <EOL> def __init__ ( self , action_noise = <NUM_LIT:0.0> , file_path = None , template_args = None ) : <EOL> if file_path is None : <EOL> if self . __class__ . FILE is None : <EOL> raise "<STR_LIT>" <EOL> file_path = osp . join ( MODEL_DIR , self . __class__ . FILE ) <EOL> if file_path . endswith ( "<STR_LIT>" ) : <EOL> lookup = mako . lookup . TemplateLookup ( directories = [ MODEL_DIR ] ) <EOL> with open ( file_path ) as template_file : <EOL> template = mako . template . Template ( <EOL> template_file . read ( ) , lookup = lookup ) <EOL> content = template . render ( <EOL> opts = template_args if template_args is not None else { } , <EOL> ) <EOL> tmp_f , file_path = tempfile . mkstemp ( text = True ) <EOL> with open ( file_path , '<STR_LIT:w>' ) as f : <EOL> f . write ( content ) <EOL> self . model = MjModel ( file_path ) <EOL> os . close ( tmp_f ) <EOL> else : <EOL> self . model = MjModel ( file_path ) <EOL> self . data = self . model . data <EOL> self . viewer = None <EOL> self . init_qpos = self . model . data . qpos <EOL> self . init_qvel = self . model . data . qvel <EOL> self . init_qacc = self . model . data . qacc <EOL> self . init_ctrl = self . model . data . ctrl <EOL> self . qpos_dim = self . init_qpos . size <EOL> self . qvel_dim = self . init_qvel . size <EOL> self . ctrl_dim = self . init_ctrl . size <EOL> self . action_noise = action_noise <EOL> if "<STR_LIT>" in self . model . numeric_names : <EOL> frame_skip_id = self . model . numeric_names . index ( "<STR_LIT>" ) <EOL> addr = self . model . numeric_adr . flat [ frame_skip_id ] <EOL> self . frame_skip = int ( self . model . numeric_data . flat [ addr ] ) <EOL> else : <EOL> self . frame_skip = <NUM_LIT:1> <EOL> if "<STR_LIT>" in self . model . numeric_names : <EOL> init_qpos_id = self . model . numeric_names . index ( "<STR_LIT>" ) <EOL> addr = self . model . numeric_adr . flat [ init_qpos_id ] <EOL> size = self . model . numeric_size . flat [ init_qpos_id ] <EOL> init_qpos = self . model . numeric_data . flat [ addr : addr + size ] <EOL> self . init_qpos = init_qpos <EOL> self . dcom = None <EOL> self . current_com = None <EOL> self . reset ( ) <EOL> super ( MujocoEnv , self ) . __init__ ( ) <EOL> @ property <EOL> @ overrides <EOL> def action_space ( self ) : <EOL> bounds = self . model . actuator_ctrlrange <EOL> lb = bounds [ : , <NUM_LIT:0> ] <EOL> ub = bounds [ : , <NUM_LIT:1> ] <EOL> return spaces . Box ( lb , ub ) <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> @ property <EOL> def action_bounds ( self ) : <EOL> return self . action_space . bounds <EOL> def reset_mujoco ( self ) : <EOL> self . model . data . qpos = self . init_qpos + np . random . normal ( size = self . init_qpos . shape ) * <NUM_LIT> <EOL> self . model . data . qvel = self . init_qvel + np . random . normal ( size = self . init_qvel . shape ) * <NUM_LIT:0.1> <EOL> self . model . data . qacc = self . init_qacc <EOL> self . model . data . ctrl = self . init_ctrl <EOL> @ overrides <EOL> def reset ( self ) : <EOL> self . reset_mujoco ( ) <EOL> self . model . forward ( ) <EOL> self . current_com = self . model . data . com_subtree [ <NUM_LIT:0> ] <EOL> self . dcom = np . zeros_like ( self . current_com ) <EOL> return self . get_current_obs ( ) <EOL> def get_current_obs ( self ) : <EOL> return self . _get_full_obs ( ) <EOL> def _get_full_obs ( self ) : <EOL> data = self . model . data <EOL> cdists = np . copy ( self . model . geom_margin ) . flat <EOL> for c in self . model . data . contact : <EOL> cdists [ c . geom2 ] = min ( cdists [ c . geom2 ] , c . dist ) <EOL> obs = np . concatenate ( [ <EOL> data . qpos . flat , <EOL> data . qvel . flat , <EOL> data . cinert . flat , <EOL> data . cvel . flat , <EOL> data . qfrc_actuator . flat , <EOL> data . cfrc_ext . flat , <EOL> data . qfrc_constraint . flat , <EOL> cdists , <EOL> self . dcom . flat , <EOL> ] ) <EOL> return obs <EOL> @ property <EOL> def _state ( self ) : <EOL> return np . concatenate ( [ <EOL> self . model . data . qpos . flat , <EOL> self . model . data . qvel . flat <EOL> ] ) <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 forward_dynamics ( self , action ) : <EOL> self . model . data . ctrl = self . inject_action_noise ( action ) <EOL> for _ in range ( self . frame_skip ) : <EOL> self . model . step ( ) <EOL> self . model . forward ( ) <EOL> new_com = self . model . data . com_subtree [ <NUM_LIT:0> ] <EOL> self . dcom = new_com - self . current_com <EOL> self . current_com = new_com <EOL> def get_viewer ( self ) : <EOL> if self . viewer is None : <EOL> self . viewer = MjViewer ( ) <EOL> self . viewer . start ( ) <EOL> self . viewer . set_model ( self . model ) <EOL> return self . viewer <EOL> def render ( self ) : <EOL> viewer = self . get_viewer ( ) <EOL> viewer . loop_once ( ) <EOL> def start_viewer ( self ) : <EOL> viewer = self . get_viewer ( ) <EOL> if not viewer . running : <EOL> viewer . start ( ) <EOL> def stop_viewer ( self ) : <EOL> if self . viewer : <EOL> self . viewer . finish ( ) <EOL> def release ( self ) : <EOL> from rllab . mujoco_py . mjlib import mjlib <EOL> mjlib . mj_deleteModel ( self . model . _wrapped ) <EOL> mjlib . mj_deleteData ( self . data . _wrapped ) <EOL> def get_body_xmat ( self , body_name ) : <EOL> idx = self . model . body_names . index ( body_name ) <EOL> return self . model . data . xmat [ idx ] . reshape ( ( <NUM_LIT:3> , <NUM_LIT:3> ) ) <EOL> def get_body_com ( self , body_name ) : <EOL> idx = self . model . body_names . index ( body_name ) <EOL> return self . model . data . com_subtree [ idx ] <EOL> def get_body_comvel ( self , body_name ) : <EOL> idx = self . model . body_names . index ( body_name ) <EOL> return self . model . body_comvels [ idx ] <EOL> def print_stats ( self ) : <EOL> super ( MujocoEnv , self ) . print_stats ( ) <EOL> print "<STR_LIT>" % len ( self . model . data . qpos ) <EOL> def action_from_key ( self , key ) : <EOL> raise NotImplementedError </s>
