def main(unused_argv): args = flags.FLAGS # checks output dir and log file out_dir = os.path.join(args.output, get_file_name_without_extension(args.replays)) create_clear_dir(out_dir, args.clear) change_log_handler(os.path.join(out_dir, 'data_collection.log'), args.verbosity) # save args with open(os.path.join(out_dir, 'args.json'), 'w') as fp: json.dump({k: args[k].value for k in args}, fp, indent=4) # create environment and collect all data from the replay # (following AIF definitions in "sc2scenarios.scenarios.assault.spaces.caml_year1_eval.get_agent_interface_format") logging.info( 'Collecting all environment data from replay file: "{}"...'.format( args.replays)) env = SC2Environment(args.replays, args.step_mul, args.replay_sc2_version, feature_screen=(192, 144), feature_minimap=72, camera_width_world_units=48, crop_to_playable_area=args.crop_to_playable_area, use_raw_units=True, action_space=ActionSpace.RAW) env_data = env.collect_all_data() agent_obs = env_data.observations # create agent according to arguments if args.task_module is not None and \ args.policy is not None and \ args.environment is not None: from interestingness_xdrl.agents.sc2_caml_y1 import SC2CAMLY1Agent logging.info('Loading reaver agent...') agent = SC2CAMLY1Agent(env.agent_interface_format, args.seed, args.sc2data_root, args.task_module, args.environment, args.policy) logging.info('Collecting interaction data for {} steps...'.format( len(agent_obs))) dataset = agent.get_interaction_datapoints(agent_obs, args.batch_size) else: logging.info('Could not determine agent to load, missing arguments...') return # saves data data_file = os.path.join(out_dir, 'interaction_data.pkl.gz') logging.info('Saving results to\n\t"{}"...'.format(data_file)) save_object(dataset, data_file) logging.info('Finished after {} timesteps ({} episodes)!'.format( len(dataset), len(np.where([dp.new_episode for dp in dataset])[0])))
def main(unused_args): args = flags.FLAGS # check for mac OS if platform.system() != 'Darwin': raise ValueError( 'Highlights extraction is currently not supported in non-macOS platforms.' ) # checks output dir and log file out_dir = args.output create_clear_dir(out_dir, args.clear) change_log_handler(os.path.join(out_dir, 'video.log'), args.verbosity) logging.info( '===================================================================') # save args with open(os.path.join(out_dir, 'args.json'), 'w') as fp: json.dump({k: args[k].value for k in args}, fp, indent=4) # collect images env = SC2Environment(args.replays, args.step_mul, args.replay_sc2_version, 1, args.window_size, args.hide_hud, True) env_data = env.collect_all_data() # saves single video or per episode # TODO more replays replay_name = get_file_name_without_extension(args.replays) frame_buffers = {} if args.separate: new_eps = list(env_data.new_episodes) new_eps.append(len(env_data.frames)) frame_buffers.update({ os.path.join(out_dir, '{}-{}.mp4'.format(replay_name, i)): env_data.frames[new_eps[i]:new_eps[i + 1]] for i in range(len(new_eps) - 1) }) else: frame_buffers[os.path.join( out_dir, '{}.mp4'.format(replay_name))] = env_data.frames for video_file, frames in frame_buffers.items(): logging.info('Got {} video frames, saving to {}...'.format( len(frames), video_file)) save_video(frames, video_file, args.fps, args.crf) logging.info('Done!')
def main(unused_argv): args = flags.FLAGS # check for mac OS if platform.system() != 'Darwin': raise ValueError('Highlights extraction is currently not supported in non-macOS platforms.') # checks output dir and log file out_dir = os.path.join(args.output, get_file_name_without_extension(args.replays)) create_clear_dir(out_dir, args.clear) change_log_handler(os.path.join(out_dir, 'highlights.log'), args.verbosity) logging.info('===================================================================') # save args with open(os.path.join(out_dir, 'args.json'), 'w') as fp: json.dump({k: args[k].value for k in args}, fp, indent=4) # load data file if not os.path.isfile(args.data): raise ValueError('Could not find interaction data file in {}'.format(args.data)) interaction_data = load_object(args.data) logging.info('Loaded interaction data corresponding to {} timesteps from: {}'.format( len(interaction_data), args.data)) # load full analysis if not os.path.isfile(args.analysis): raise ValueError('Could not find full analysis data file in {}'.format(args.analysis)) analyses = FullAnalysis.load(args.analysis, interaction_data) logging.info('Loaded full analysis data file from: {}'.format(args.analysis)) logging.info('___________________________________________________________________') logging.info('Collecting visual information from SC2 by replaying \'{}\'...'.format(args.replays)) # collect images env = SC2Environment( args.replays, RECORD_STEP_MUL, args.replay_sc2_version, 1, args.window_size, args.hide_hud, True) env_data = env.collect_all_data() # collects and saves highlights explainer = HighlightsReport(analyses, out_dir, env_data.frames, args.step_mul / RECORD_STEP_MUL, args.record_time_steps, args.max_highlights_per_elem, args.fps, args.fade_ratio) explainer.create() logging.info('Finished after {} timesteps!'.format(len(env_data.frames)))
def process_replay(args, file, vae_model, session): logging.info('Replaying {}...'.format(file)) # create environment env = SC2Environment(file, args.step_mul, 1., args.replay_sc2_version, window_size=args.window_size, hide_hud=args.hide_hud, capture_screen=args.video_eps > 0, feature_screen=args.obs_spatial_dim, use_feature_units=True) # collect all data from the replay env_data = env.collect_all_data(args.max_eps) screen_buffer = env_data.frames agent_obs = env_data.observations agent_actions = env_data.actions max_video_step = env_data.new_episodes[args.video_eps] out_dir = os.path.join(args.output, get_file_name_without_extension(file)) create_clear_dir(out_dir, args.clear) if args.video_eps > 0: file_path = os.path.join( out_dir, '{}-screen.mp4'.format(get_file_name_without_extension(file))) save_video(screen_buffer[:max_video_step], file_path, args.fps, args.crf) # create agent agent = SC2ReaverAgent( env.agent_interface_format, args.seed, args.results_dir, args.experiment, args.env, args.obs_features, args.action_set, args.action_spatial_dim, args.safety, args.safe_distance, args.safe_flee_angle_range, args.safe_margin, args.safe_unsafe_steps, args.gin_files, args.gin_bindings, args.agent, args.gpu) # create converter from/to agent observations using VAE converter = ReaverVAESampleConverter(env.agent_interface_format, vae_model, args.action_set, args.action_spatial_dim, args.det_vae) # collects real observation images and saves video real_ag_obs_img_buffer = converter.to_images(agent_obs[:max_video_step]) file_path = os.path.join( out_dir, '{}-orig-layers.mp4'.format(get_file_name_without_extension(file))) save_video(real_ag_obs_img_buffer, file_path, args.fps, args.crf, verify_input=True, color='white') # gets predictive rollouts of length l for each step and using the agent's policy pe_model = BNN.load(args.pe_model, converter.observation_dim, converter.action_dim, RWD_DIM, session) evals = [] for length in args.rollout_lengths: out_dir_ = os.path.join(out_dir, str(length)) create_clear_dir(out_dir_, args.clear) logging.info( 'Evaluating PE for rollouts of length {}, saving results to\n\t"{}"' .format(length, out_dir_)) obs_rollouts, rwd_rollouts = get_agent_rollouts( pe_model, agent, converter, agent_obs, agent_actions, RWD_DIM, args.batch_size, length, args.det_pe) # converts last latent observation as predicted by rollouts of each predictor to an agent observation, using vae rollout_rewards = [ rwd_rollouts[-1, 0, i, :, :] for i in range(pe_model.num_nets) ] rollout_agent_obs = [ converter.to_agent_observations( obs_rollouts[-1, 0, i, :, :], rollout_rewards[i], agent_obs[length:] + agent_obs[-length:]) for i in range(pe_model.num_nets) ] # gets time-aligned predicted observations pred_agent_obs = [ agent_obs[:length] + rollout_agent_obs[k][:-length] for k in range(pe_model.num_nets) ] # calculate mean obs diff at each timestep for each predictor real_datapoints = agent.get_interaction_datapoints(agent_obs) pred_datapoints = [ agent.get_interaction_datapoints(pred_agent_obs[k]) for k in range(pe_model.num_nets) ] mean_diffs = np.array([[ np.mean( list( get_observation_differences( real_datapoints[i].observation, pred_datapoints[k] [i].observation)[3]['feature_screen'].values())) for i in range(len(real_datapoints)) ] for k in range(pe_model.num_nets)]) # selects best predictor (lower mean diff) at each time step, save video best_idxs = np.argmin(mean_diffs, axis=0) best_pred_ag_obs = [ pred_agent_obs[best_idxs[i]][i] for i in range(len(best_idxs)) ] pred_ag_obs_img_buffer = converter.to_images( best_pred_ag_obs[:max_video_step]) file_path = os.path.join( out_dir_, '{}-best-pred-layers.mp4'.format( get_file_name_without_extension(file))) save_video(pred_ag_obs_img_buffer, file_path, args.fps, args.crf, verify_input=True, color='white') # collects predicted observations images (for each predictor) and saves video pred_ag_obs_img_buffers = [] for k in range(pe_model.num_nets): # first 'rollout-length' observations cannot be predicted, so just copy original ones pred_ag_obs_img_buffer = converter.to_images( pred_agent_obs[k][:max_video_step]) file_path = os.path.join( out_dir_, '{}-pred-layers-net-{}.mp4'.format( get_file_name_without_extension(file), k)) save_video(pred_ag_obs_img_buffer, file_path, args.fps, args.crf, verify_input=True, color='white') pred_ag_obs_img_buffers.append(pred_ag_obs_img_buffer) # saves the mean and variance between layers pred_ag_obs_img_buffer = [ get_mean_image([ pred_ag_obs_img_buffers[k][i] for k in range(pe_model.num_nets) ], canvas_color='white') for i in range(len(pred_ag_obs_img_buffers[0])) ] file_path = os.path.join( out_dir_, '{}-mean-pred-layers.mp4'.format( get_file_name_without_extension(file))) save_video(pred_ag_obs_img_buffer, file_path, args.fps, args.crf, verify_input=True, color='white') pred_ag_obs_img_buffer = [ get_variance_heatmap([ pred_ag_obs_img_buffers[k][i] for k in range(pe_model.num_nets) ], False, True, canvas_color='white') for i in range(len(pred_ag_obs_img_buffers[0])) ] file_path = os.path.join( out_dir_, '{}-var-pred-layers.mp4'.format( get_file_name_without_extension(file))) save_video(pred_ag_obs_img_buffer, file_path, args.fps, args.crf, verify_input=True, color='white') # evaluates real vs predicted performance for each predictor length_evals = [] for k in range(pe_model.num_nets): out_dir__ = os.path.join(out_dir_, 'net-{}-eval'.format(k)) create_clear_dir(out_dir__, args.clear) length_evals.append( compare_datapoints( real_datapoints, pred_datapoints[k], out_dir__, 'Real', 'Predictor {} ($\\mathbf{{h={}}}$)'.format(k, length), [c.name for c in vae_model.components])) evals.append(length_evals) # gets mean differences and errors from evaluation data rollout_diffs = {} for i, length_evals in enumerate(evals): # for each rollout nets_data = {} for net_evals in length_evals: # for each network in the ensemble # collect eval data for each network and eval criterion for eval_name in net_evals.keys(): eval_data = net_evals[eval_name] if eval_name not in nets_data: nets_data[eval_name] = [] nets_data[eval_name].append(eval_data) # gets mean and standard error (across networks in ensemble) for each eval criterion if i == 0: rollout_diffs = {eval_name: {} for eval_name in nets_data.keys()} for eval_name, eval_data in nets_data.items(): rollout_diffs[eval_name][str(args.rollout_lengths[i])] = [ np.mean(eval_data), np.std(eval_data) / len(eval_data) ] # saves data per rollout out_dir = os.path.join(out_dir, 'mean-eval') create_clear_dir(out_dir, args.clear) for eval_name, diffs in rollout_diffs.items(): plot_bar(diffs, 'Mean {}'.format(eval_name), os.path.join( out_dir, 'mean-{}.pdf'.format(eval_name.lower().replace( ' ', '-').replace('.', ''))), plot_mean=False, x_label='Rollout Length')
def process_replay(args, file, vae_model): logging.info('Replaying {}...'.format(file)) # create environment env = SC2Environment(file, args.step_mul, 1., args.replay_sc2_version, window_size=args.window_size, hide_hud=args.hide_hud, capture_screen=args.video_eps > 0, feature_screen=args.obs_spatial_dim, use_feature_units=True) # collect all data from the replay env_data = env.collect_all_data(args.max_eps) screen_buffer = env_data.frames agent_obs = env_data.observations agent_actions = env_data.actions max_video_step = env_data.new_episodes[args.video_eps] out_dir = os.path.join(args.output, get_file_name_without_extension(file)) create_clear_dir(out_dir, args.clear) if args.video_eps > 0: file_path = os.path.join( out_dir, '{}-screen.mp4'.format(get_file_name_without_extension(file))) save_video(screen_buffer[:max_video_step], file_path, args.fps, args.crf) # create real agent real_agent = SC2ReaverAgent( env.agent_interface_format, args.seed, args.results_dir, args.experiment, args.env, args.obs_features, args.action_set, args.action_spatial_dim, args.safety, args.safe_distance, args.safe_flee_angle_range, args.safe_margin, args.safe_unsafe_steps, args.gin_files, args.gin_bindings, args.agent, args.gpu) # create vae-based agent converter = ReaverVAESampleConverter(env.agent_interface_format, vae_model, args.action_set, args.action_spatial_dim) vae_agent = SC2ReaverVAEAgent(real_agent, converter) # collect datapoints for real and vae-reconstructed observations real_datapoints = real_agent.get_interaction_datapoints(agent_obs) vae_datapoints = vae_agent.get_interaction_datapoints( (agent_obs, agent_actions)) # collects real vs vae-reconstructed observation images and saves video real_obs_img_buffer = converter.to_images(agent_obs[:max_video_step]) file_path = os.path.join( out_dir, '{}-orig-layers.mp4'.format(get_file_name_without_extension(file))) save_video(real_obs_img_buffer, file_path, args.fps, args.crf, verify_input=True, color='white') vae_obs_img_buffer = converter.to_images( vae_agent.agent_observations[:max_video_step]) file_path = os.path.join( out_dir, '{}-vae-layers.mp4'.format(get_file_name_without_extension(file))) save_video(vae_obs_img_buffer, file_path, args.fps, args.crf, verify_input=True, color='white') # evaluates real vs VAE performance compare_datapoints(real_datapoints, vae_datapoints, out_dir, 'Real', 'VAE', [c.name for c in vae_model.components]) # mean vae variance mean_vars = np.exp(vae_agent.latent_log_vars).mean(axis=0, keepdims=True).reshape( 1, -1) plot_evolution(mean_vars, [''], 'Mean VAE Variance $z_{\sigma^2}$', output_img=os.path.join(out_dir, 'eval-vae-var.pdf'), x_label='Time') logging.info('Finished after {} timesteps!'.format(len(screen_buffer)))
def main(unused_argv): args = flags.FLAGS # check for mac OS if platform.system() != 'Darwin': raise ValueError( 'Counterfactual extraction is currently supported only in non-macOS platform.' ) # checks output dir and log file out_dir = os.path.join(args.output, get_file_name_without_extension(args.replays)) create_clear_dir(out_dir, args.clear) change_log_handler(os.path.join(out_dir, 'counterfactuals.log'), args.verbosity) # save args with open(os.path.join(out_dir, 'args.json'), 'w') as fp: json.dump({k: args[k].value for k in args}, fp, indent=4) # load data file if not os.path.isfile(args.data): raise ValueError('Could not find interaction data file in {}'.format( args.data)) interaction_data = load_object(args.data) logging.info( 'Loaded interaction data corresponding to {} timesteps from: {}'. format(len(interaction_data), args.data)) # load full analysis if not os.path.isfile(args.analysis): raise ValueError('Could not find full analysis data file in {}'.format( args.analysis)) analyses = FullAnalysis.load(args.analysis, interaction_data) logging.info('Loaded full analysis data file from: {}'.format( args.analysis)) # create environment env = SC2Environment(args.replays, args.step_mul, 1., args.replay_sc2_version, 1, True, args.window_size, args.hide_hud, True, args.obs_spatial_dim, use_feature_units=True) env_data = env.collect_all_data() # create agent according to arguments if args.experiment is not None and \ args.env is not None: from interestingness_xdrl.agents.sc2_reaver import SC2ReaverAgent logging.info('Loading reaver agent...') # create agent agent = SC2ReaverAgent(env.agent_interface_format, args.seed, args.results_dir, args.experiment, args.env, args.obs_features, args.action_set, args.action_spatial_dim, args.safety, args.safe_distance, args.safe_flee_angle_range, args.safe_margin, args.safe_unsafe_steps, args.gin_files, args.gin_bindings, args.agent, args.gpu) elif args.vae_model is not None: from imago.models.sequential.pets.converters.rb_vae_converter import RBVAESampleConverter from imago.models.behav.rb_perturb import RBPerturbModel from imago.models.sequential.pets.bnn import setup_tf from interestingness_xdrl.agents.sc2_rb_vae_bnn import SC2RBVAEBNNAgent logging.info( 'Loading agent with predictive model ensemble and reaver-like behavior via VAE...' ) # setup TF session, has_gpu = setup_tf() # create converter from/to agent observations using VAE rb_perturb_model = RBPerturbModel(args.vae_model, 'GPU:0' if has_gpu else 'cpu') converter = RBVAESampleConverter(env.agent_interface_format, rb_perturb_model, args.action_set, args.action_spatial_dim, args.det_vae) agent = SC2RBVAEBNNAgent(converter, None, args.seed) else: logging.info('Could not determine agent to load, missing arguments...') return # collects and saves counterfactuals explainer = CounterfactualsReport(analyses, out_dir, env_data.frames, agent) explainer.create() logging.info('Finished after {} timesteps!'.format(len(env_data.frames)))
def main(unused_argv): args = flags.FLAGS # checks output dir and log file out_dir = os.path.join(args.output, get_file_name_without_extension(args.replays)) create_clear_dir(out_dir, args.clear) change_log_handler(os.path.join(out_dir, 'data_collection.log'), args.verbosity) # save args with open(os.path.join(out_dir, 'args.json'), 'w') as fp: json.dump({k: args[k].value for k in args}, fp, indent=4) # create environment and collect all data from the replay logging.info('Collecting all environment data from replay file: "{}"...'.format(args.replays)) env = SC2Environment(args.replays, args.step_mul, args.replay_sc2_version, feature_screen=args.obs_spatial_dim, feature_minimap=args.obs_spatial_dim, use_feature_units=True) env_data = env.collect_all_data() agent_obs = env_data.observations agent_actions = env_data.actions # create agent according to arguments if args.experiment is not None and \ args.env is not None and \ args.vae_model is not None and \ args.pe_model is not None: from imago.models.sequential.pets.bnn import setup_tf, BNN from imago.models.sequential.pets.converters.reaver_vae_converter import ReaverVAESampleConverter from imago.models.semframe import model_remapper from interestingness_xdrl.agents.sc2_reaver_vae_bnn import SC2ReaverVAEBNNAgent logging.info('Loading reaver agent with predictive model ensemble via VAE...') # setup TF session, has_gpu = setup_tf() # loads VAE rb_perturb_model = model_remapper.load_model(args.vae_model, device=None if has_gpu else 'cpu') # create converter from/to agent observations using VAE converter = ReaverVAESampleConverter(env.agent_interface_format, rb_perturb_model, args.action_set, args.action_spatial_dim, args.det_vae) # loads probabilistic ensemble of predictive models pe_model = BNN.load(args.pe_model, converter.observation_dim, converter.action_dim, RWD_DIM, session) agent = SC2ReaverVAEBNNAgent( converter, pe_model, env.agent_interface_format, args.seed, args.results_dir, args.experiment, args.env, args.obs_features, args.action_set, args.action_spatial_dim, args.safety, args.safe_distance, args.safe_flee_angle_range, args.safe_margin, args.safe_unsafe_steps, args.gin_files, args.gin_bindings, args.agent, args.gpu, args.horizon, args.det_pe) logging.info('Collecting interaction data for {} steps...'.format(len(agent_obs))) dataset = agent.get_interaction_datapoints((agent_obs, agent_actions)) elif args.vae_model is not None and args.pe_model is not None: from imago.models.sequential.pets.converters.rb_vae_converter import RBVAESampleConverter from imago.models.behav.rb_perturb import RBPerturbModel from imago.models.sequential.pets.bnn import setup_tf, BNN from interestingness_xdrl.agents.sc2_rb_vae_bnn import SC2RBVAEBNNAgent logging.info('Loading agent with predictive model ensemble and reaver-like behavior via VAE...') # setup TF session, has_gpu = setup_tf() # loads VAE rb_perturb_model = RBPerturbModel(args.vae_model, 'GPU:0' if has_gpu else 'cpu') # create converter from/to agent observations using VAE converter = RBVAESampleConverter(env.agent_interface_format, rb_perturb_model, args.action_set, args.action_spatial_dim, args.det_vae) # loads probabilistic ensemble of predictive models pe_model = BNN.load(args.pe_model, converter.observation_dim, converter.action_dim, RWD_DIM, session) agent = SC2RBVAEBNNAgent(converter, pe_model, args.seed, args.horizon, args.det_pe) logging.info('Collecting interaction data for {} steps...'.format(len(agent_obs))) dataset = agent.get_interaction_datapoints(agent_obs) elif args.experiment is not None and args.env is not None: from interestingness_xdrl.agents.sc2_reaver import SC2ReaverAgent logging.info('Loading reaver agent...') agent = SC2ReaverAgent( env.agent_interface_format, args.seed, args.results_dir, args.experiment, args.env, args.obs_features, args.action_set, args.action_spatial_dim, args.safety, args.safe_distance, args.safe_flee_angle_range, args.safe_margin, args.safe_unsafe_steps, args.gin_files, args.gin_bindings, args.agent, args.gpu) logging.info('Collecting interaction data for {} steps...'.format(len(agent_obs))) dataset = agent.get_interaction_datapoints(agent_obs) else: logging.info('Could not determine agent to load, missing arguments...') return # saves data data_file = os.path.join(out_dir, 'interaction_data.pkl.gz') save_object(dataset, data_file) logging.info('Finished after {} timesteps, saved results to\n\t"{}"'.format(len(dataset), data_file))
def main(unused_argv): args = flags.FLAGS # checks output dir and files create_clear_dir(args.output, args.clear) change_log_handler(os.path.join(args.output, 'select.log'), 1) # gets files' properties files = get_files_with_extension(args.replays, REPLAY_EXT) logging.info('=============================================') logging.info('Found {} replays in {}'.format(len(files), args.replays)) files_attrs = [] for file in files: date = os.path.getmtime(file) size = os.path.getsize(file) files_attrs.append((file, date, size)) # sort and gets files' sizes files_attrs.sort(key=lambda f: f[1]) sizes = np.array([file_attrs[2] for file_attrs in files_attrs]) plot_evolution(sizes.reshape(1, -1), [''], 'Replay File Size', output_img=os.path.join(args.output, 'files-size.pdf'), y_label='bytes') # selects and copies files size_diffs = sizes[1:] - sizes[:-1] size_diff_thresh = -(np.mean(np.abs(size_diffs)) + THRESH_STDS * np.std(np.abs(size_diffs))) max_size_idxs = np.where(size_diffs < size_diff_thresh)[0].tolist() max_size_idxs.append(len(sizes) - 1) # add last replay idx logging.info('{} replay files selected:'.format(len(max_size_idxs))) total_steps = 0 for i in range(len(max_size_idxs)): idx = max_size_idxs[i] file, date, size = files_attrs[idx] logging.info( 'Copying "{}" (Created: {:%m-%d %H:%M:%S}, Size: {}b) to "{}"'. format(file, datetime.datetime.fromtimestamp(date), size, args.output)) out_file = os.path.join(args.output, f'eps_{i}.{REPLAY_EXT}') shutil.copy(file, out_file) # check also timesteps metadata file from CAML scenarios timesteps_file = os.path.join( os.path.dirname(file), '{}_timesteps.json'.format(get_file_name_without_extension(file))) if os.path.isfile(timesteps_file): # collect all timesteps info from the files json_steps = [] for j in range(0 if i == 0 else max_size_idxs[i - 1], idx + 1): f = files_attrs[j][0] with open( os.path.join( os.path.dirname(f), '{}_timesteps.json'.format( get_file_name_without_extension(f))), 'r') as fp: json_steps.extend(json.load(fp)) # save timesteps file with all steps info with open(os.path.join(args.output, f'eps_{i}_timesteps.json'), 'w') as fp: json.dump(json_steps, fp) total_steps += len(json_steps) if total_steps > 0: logging.info( '{} total steps selected from replays.'.format(total_steps)) # runs replays to get stats if args.analyze: steps_queue = mp.JoinableQueue() sample_processor = _ReplayAnalyzerProcessor(steps_queue, args.step_mul) replayer_runner = ReplayProcessRunner(args.output, sample_processor, args.replay_sc2_version, mp.cpu_count(), player_ids=PLAYER_PERSPECTIVE) replayer_runner.run() steps_queue.put(None) # process results steps = [] while True: ep_steps = steps_queue.get() if ep_steps is None: break steps.extend(ep_steps) time.sleep(0.01) steps = np.array([steps]) plot_evolution(steps, [''], 'Steps per Episode', output_img=os.path.join(args.output, 'ep-steps.pdf'), y_label='Num. Steps') logging.info('=============================================') logging.info('Got {} episodes, {} total steps, mean: {}'.format( steps.shape[1], steps.sum(), steps.mean())) logging.info('Finished')