def main()

in train_procgen/graph.py [0:0]


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--distribution_mode', type=str, default='easy', help="Environment distribution_mode ('easy' or 'hard')")
    parser.add_argument('--normalize_and_reduce', dest='normalize_and_reduce', action='store_true')
    parser.add_argument('--restrict_training_set', dest='restrict_training_set', action='store_true')
    parser.add_argument('--save', dest='save', action='store_true')
    args = parser.parse_args()

    run_directory_prefix = main_pcg_sample_entry(args.distribution_mode, args.normalize_and_reduce, args.restrict_training_set)

    plt.tight_layout()

    if args.save:
        suffix = '-mean' if args.normalize_and_reduce else ''
        plt.savefig(f'results/{run_directory_prefix}{suffix}.pdf')
    else:
        plt.show()