<s> from rllab . algos . base import RLAlgorithm <EOL> import theano . tensor as TT <EOL> import numpy as np <EOL> from rllab . misc import ext <EOL> from rllab . misc . special import discount_cumsum <EOL> from rllab . sampler import parallel_sampler , stateful_pool <EOL> from rllab . sampler . utils import rollout <EOL> from rllab . core . serializable import Serializable <EOL> import rllab . misc . logger as logger <EOL> import rllab . plotter as plotter <EOL> import cma_es_lib <EOL> def sample_return ( G , params , max_path_length , discount ) : <EOL> G . policy . set_param_values ( params ) <EOL> path = rollout ( <EOL> G . env , <EOL> G . policy , <EOL> max_path_length , <EOL> ) <EOL> path [ "<STR_LIT>" ] = discount_cumsum ( path [ "<STR_LIT>" ] , discount ) <EOL> path [ "<STR_LIT>" ] = sum ( path [ "<STR_LIT>" ] ) <EOL> return path <EOL> class CMAES ( RLAlgorithm , Serializable ) : <EOL> def __init__ ( <EOL> self , <EOL> env , <EOL> policy , <EOL> n_itr = <NUM_LIT> , <EOL> max_path_length = <NUM_LIT> , <EOL> discount = <NUM_LIT> , <EOL> sigma0 = <NUM_LIT:1.> , <EOL> batch_size = None , <EOL> plot = False , <EOL> ** kwargs <EOL> ) : <EOL> """<STR_LIT>""" <EOL> Serializable . quick_init ( self , locals ( ) ) <EOL> self . env = env <EOL> self . policy = policy <EOL> self . plot = plot <EOL> self . sigma0 = sigma0 <EOL> self . discount = discount <EOL> self . max_path_length = max_path_length <EOL> self . n_itr = n_itr <EOL> self . batch_size = batch_size <EOL> def train ( self ) : <EOL> cur_std = self . sigma0 <EOL> cur_mean = self . policy . get_param_values ( ) <EOL> es = cma_es_lib . CMAEvolutionStrategy ( <EOL> cur_mean , cur_std ) <EOL> parallel_sampler . populate_task ( self . env , self . policy ) <EOL> if self . plot : <EOL> plotter . init_plot ( self . env , self . policy ) <EOL> cur_std = self . sigma0 <EOL> cur_mean = self . policy . get_param_values ( ) <EOL> itr = <NUM_LIT:0> <EOL> while itr < self . n_itr and not es . stop ( ) : <EOL> if self . batch_size is None : <EOL> xs = es . ask ( ) <EOL> xs = np . asarray ( xs ) <EOL> infos = ( <EOL> stateful_pool . singleton_pool . run_map ( sample_return , [ ( x , self . max_path_length , <EOL> self . discount ) for x in xs ] ) ) <EOL> else : <EOL> cum_len = <NUM_LIT:0> <EOL> infos = [ ] <EOL> xss = [ ] <EOL> done = False <EOL> while not done : <EOL> sbs = stateful_pool . singleton_pool . n_parallel * <NUM_LIT:2> <EOL> xs = es . ask ( sbs ) <EOL> xs = np . asarray ( xs ) <EOL> xss . append ( xs ) <EOL> sinfos = stateful_pool . singleton_pool . run_map ( <EOL> sample_return , [ ( x , self . max_path_length , self . discount ) for x in xs ] ) <EOL> for info in sinfos : <EOL> infos . append ( info ) <EOL> cum_len += len ( info [ '<STR_LIT>' ] ) <EOL> if cum_len >= self . batch_size : <EOL> xs = np . concatenate ( xss ) <EOL> done = True <EOL> break <EOL> fs = - np . array ( [ info [ '<STR_LIT>' ] [ <NUM_LIT:0> ] for info in infos ] ) <EOL> xs = xs [ : len ( fs ) ] <EOL> es . tell ( xs , fs ) <EOL> logger . push_prefix ( '<STR_LIT>' % itr ) <EOL> logger . record_tabular ( '<STR_LIT>' , itr ) <EOL> logger . record_tabular ( '<STR_LIT>' , np . mean ( cur_std ) ) <EOL> undiscounted_returns = np . array ( <EOL> [ info [ '<STR_LIT>' ] for info in infos ] ) <EOL> logger . record_tabular ( '<STR_LIT>' , <EOL> np . mean ( undiscounted_returns ) ) <EOL> logger . record_tabular ( '<STR_LIT>' , <EOL> np . mean ( undiscounted_returns ) ) <EOL> logger . record_tabular ( '<STR_LIT>' , <EOL> np . max ( undiscounted_returns ) ) <EOL> logger . record_tabular ( '<STR_LIT>' , <EOL> np . min ( undiscounted_returns ) ) <EOL> logger . record_tabular ( '<STR_LIT>' , <EOL> np . mean ( fs ) ) <EOL> logger . record_tabular ( '<STR_LIT>' , <EOL> np . mean ( [ len ( info [ '<STR_LIT>' ] ) for info in infos ] ) ) <EOL> self . env . log_diagnostics ( infos ) <EOL> self . policy . log_diagnostics ( infos ) <EOL> logger . save_itr_params ( itr , dict ( <EOL> itr = itr , <EOL> policy = self . policy , <EOL> env = self . env , <EOL> ) ) <EOL> logger . dump_tabular ( with_prefix = False ) <EOL> if self . plot : <EOL> plotter . update_plot ( self . policy , self . max_path_length ) <EOL> logger . pop_prefix ( ) <EOL> itr += <NUM_LIT:1> <EOL> self . policy . set_param_values ( es . result ( ) [ <NUM_LIT:0> ] ) <EOL> parallel_sampler . terminate_task ( ) </s>
