Пример #1
0
def main(unused_argv):
    args = flags.FLAGS

    if not os.path.isfile(args.data):
        raise ValueError('Could not find interaction data file in {}'.format(args.data))

    # checks output dir and log file
    out_dir = os.path.join(args.output, get_directory_name(args.data))
    create_clear_dir(out_dir, args.clear)
    change_log_handler(os.path.join(out_dir, 'analyses.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
    interaction_data = load_object(args.data)
    num_eps = len(np.where([dp.new_episode for dp in interaction_data])[0])
    logging.info('Loaded interaction data corresponding to {} timesteps ({} episodes)from: {}'.format(
        len(interaction_data), num_eps, args.data))

    # load analysis config
    if not os.path.isfile(args.config):
        raise ValueError('Could not find analysis configuration file in {}'.format(args.config))
    config = AnalysisConfiguration.load_json(args.config)
    logging.info('Loaded analysis configuration file from: {}'.format(args.config))
    config.save_json(os.path.join(out_dir, os.path.basename(args.config)))

    # creates full analysis with all analyses
    analysis = FullAnalysis(interaction_data, config, args.img_format)
    logging.info('{} total analyses to be performed...'.format(len(analysis)))

    # runs and saves results
    analysis.analyze(out_dir)
    analysis.save(os.path.join(out_dir, 'analyses.pkl.gz'))
def main(unused_argv):
    args = flags.FLAGS

    # checks output dir and log file
    create_clear_dir(args.output, args.clear)
    change_log_handler(os.path.join(args.output, 'evaluate_pe.log'),
                       args.verbosity)

    # save args
    with open(os.path.join(args.output, 'args.json'), 'w') as fp:
        json.dump({k: args[k].value for k in args}, fp, indent=4)

    # setup TF
    session, has_gpu = setup_tf()

    # loads models
    vae_model = model_remapper.load_model(args.vae_model,
                                          device=None if has_gpu else 'cpu')

    # checks input files
    if os.path.isfile(args.replays):
        files = [args.replays]
    elif os.path.isdir(args.replays):
        files = list(get_files_with_extension(args.replays, 'SC2Replay'))
    else:
        raise ValueError(
            'Input path is not a valid file or directory: {}.'.format(
                args.input))

    # process files
    for file in files:
        process_replay(args, file, vae_model, session)
Пример #3
0
def main(unused_argv):
    args = flags.FLAGS

    # checks output dir and log file
    out_dir = os.path.join(args.output)
    create_clear_dir(out_dir, args.clear)
    change_log_handler(os.path.join(out_dir, 'dataset.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))

    # collects and saves datasets
    report = DatasetReport(analyses, out_dir)
    report.create()
    logging.info('Finished after {} timesteps!'.format(len(interaction_data)))
def main(unused_argv):
    args = flags.FLAGS

    # checks output dir and log file
    create_clear_dir(args.output, args.clear)
    change_log_handler(os.path.join(args.output, 'features.log'),
                       args.verbosity)
    logging.info(
        '===================================================================')

    # save args
    with open(os.path.join(args.output, 'args.json'), 'w') as fp:
        json.dump({k: args[k].value for k in args}, fp, indent=4)

    # 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)
    logging.info('Loaded full analysis data file from: {}'.format(
        args.analysis))

    # collects and saves features
    explainer = SC2FeatureReport(analyses, args.output, args.features,
                                 args.time_steps)
    explainer.create()
Пример #5
0
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])))
Пример #6
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!')
Пример #7
0
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 main(unused_argv):
    args = flags.FLAGS

    # checks output dir and log file
    create_clear_dir(args.output, args.clear)
    change_log_handler(os.path.join(args.output, 'tracker.log'), args.verbosity)
    logging.info('===================================================================')

    # save args
    with open(os.path.join(args.output, 'args.json'), 'w') as fp:
        json.dump({k: args[k].value for k in args}, fp, indent=4)

    logging.info('___________________________________________________________________')
    logging.info('Tracking units in SC2 by replaying \'{}\'...'.format(args.replay_file))

    env = SC2Environment(
        args.replay_file, args.step_mul, 1., args.replay_sc2_version, 1, False, args.window_size, args.hide_hud, True,
        FEATURE_DIMENSIONS, CAMERA_WIDTH, True)
    env.start()

    # gets perspective transformation matrix
    left = int(0.04 * FEATURE_DIMENSIONS)
    right = int(0.96 * FEATURE_DIMENSIONS)
    top = int(0.04 * FEATURE_DIMENSIONS)
    bottom = int(0.7 * FEATURE_DIMENSIONS)
    m = cv2.getPerspectiveTransform(
        np.float32([(left, top), (right, top), (right, bottom), (left, bottom)]),
        np.float32([(0.09, 0.05), (0.91, 0.05), (1.04, 0.76), (-0.04, 0.76)]) * np.float32(env.visual_observation.size))

    # creates "square matrix"
    border_mat = np.zeros(tuple(env.agent_interface_format.feature_dimensions.screen))
    border_mat[0, :] = border_mat[-1, :] = border_mat[:, 0] = border_mat[:, -1] = 1

    ep = -1
    video_writer = None
    replay_file = os.path.basename(args.replay_file)
    while not env.finished:

        if env.t == 0 or env.new_episode:
            ep += 1

            if video_writer is not None:
                video_writer.close()

            # creates video writer
            ext_idx = replay_file.lower().find('.sc2replay')
            output_file = os.path.join(
                args.output, '{}-{}.mp4'.format(replay_file[:ext_idx], ep))
            video_writer = skvideo.io.FFmpegWriter(
                output_file, inputdict={'-r': str(args.fps)}, outputdict={'-crf': str(args.crf), '-pix_fmt': 'yuv420p'})

            logging.info('Recording episode {} of replay \'{}\' to \'{}\'...'.format(ep, replay_file, output_file))

        # capture units
        img = env.visual_observation
        if img is not None:
            masks_colors = [(border_mat, [255, 255, 255]),
                            (env.agent_obs.observation.feature_screen.player_relative == 1, [255, 0, 0]),
                            (env.agent_obs.observation.feature_screen.player_relative == 4, [0, 0, 255])]
            for mask, color in masks_colors:
                mask = np.asarray(mask * 100, dtype=np.uint8)
                mask = cv2.warpPerspective(mask, m, img.size)
                mask = cv2.dilate(mask, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (10, 10)))
                mask = cv2.GaussianBlur(mask, (51, 51), 0)
                mask = Image.fromarray(mask)

                overlay = np.zeros((img.size[1], img.size[0], 3), dtype=np.uint8)
                for i in range(len(color)):
                    overlay[:, :, i] = color[i]
                overlay = Image.fromarray(overlay)

                img = Image.composite(overlay, img, mask)
            video_writer.writeFrame(np.array(img))

        env.step()

    env.stop()

    logging.info('Finished after {} timesteps!'.format(env.t))
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)))
Пример #10
0
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')