Douzero_Resnet/generate_eval_data.py

import argparse
import pickle
import numpy as np

deck = []
for i in range(3, 15):
    deck.extend([i for _ in range(8)])
deck.extend([17 for _ in range(8)])
deck.extend([20, 20, 30, 30])

def get_parser():
    parser = argparse.ArgumentParser(description='DouZero: random data generator')
    parser.add_argument('--output', default='eval_data', type=str)
    parser.add_argument('--num_games', default=10000, type=int)
    return parser
    
def generate():
    _deck = deck.copy()
    np.random.shuffle(_deck)
    card_play_data = {'landlord': _deck[:33],
                      'landlord_up': _deck[33:58],
                      'landlord_front': _deck[58:83],
                      'landlord_down': _deck[83:108],
                      'three_landlord_cards': _deck[25:33],
                      'three_landlord_cards_all': _deck[25:33],
                      }
    for key in card_play_data:
        card_play_data[key].sort()
    return card_play_data


if __name__ == '__main__':
    flags = get_parser().parse_args()
    output_pickle = flags.output + '.pkl'

    print("output_pickle:", output_pickle)
    print("generating data...")

    data = []
    for _ in range(flags.num_games):
        data.append(generate())

    print("saving pickle file...")
    with open(output_pickle,'wb') as g:
        pickle.dump(data,g,pickle.HIGHEST_PROTOCOL)