replication_project/SL_model.py at master · danieldritter/replication_project · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
import numpy as np
import datetime
from collections import defaultdict
import tensorflow as tf
from tensorflow.keras import Model
from tensorflow.keras.losses import categorical_crossentropy
from tensorflow.keras.optimizers import Adam
import pickle

from data import process
from SL.encoder.encoder import Encoder
from SL.decoder.decoder import Decoder
from constants.constants import SEASON, UNIT_POWER, ORDER_DICT
from AbstractActor import AbstractActor

def set_sl_weights(new_weights, sl_model, state_inputs, prev_order_inputs, prev_orders_game_labels, season_names, board_dict_list):
    # Finds winning power to use in network (looking at 1st game)
    i = 0

    last_turn = board_dict_list[i][-1]
    prov_num_dict = defaultdict(int)
    for province in last_turn:
        if "unit_power" in last_turn[province].keys():
            prov_num_dict[last_turn[province]["unit_power"]] += 1
    power = max(prov_num_dict, key=prov_num_dict.get)

    # Parsing just the season(not year)
    # Not sure about these conversions
    powers_seasons = []
    curr_board_dict = board_dict_list[i]
    # extracting seasons and powers for film
    for j in range(len(season_names[i])):
        # print(season_names[i][j][0])
        powers_seasons.append(
            SEASON[season_names[i][j][0]] + UNIT_POWER[power])
    # print(powers_seasons)

    # casting encoder and decoder inputs to floats
    powers_seasons = tf.convert_to_tensor(powers_seasons,
                                            dtype=tf.float32)
    state_input = tf.convert_to_tensor(state_inputs[i],
                                        dtype=tf.float32)
    order_inputs = tf.convert_to_tensor(prev_order_inputs[i],
                                        dtype=tf.float32)
    season_input = season_names[i]

    # applying SL model
    orders_probs, position_lists = sl_model.call(state_input,
                                             order_inputs,
                                             powers_seasons,
                                             season_input,
                                             curr_board_dict,
                                             power)
    sl_model.set_weights(new_weights)

class SL_model(AbstractActor):
    '''
    The supervised learning Actor for the Diplomacy game
    '''

    def loss(self, prev_order_phase_labels, probs, position_lists, power):
        '''
        Function to compute the loss of the SL Model

        Keyword Args:
        labels - the previous order game labels
        probs - the probability distribution output over the orders
        prev_order_phase_labels - the labels for the phases of the game
        position_lists - the list of positions that a power controlled

        Return:
        the crossentropy loss for the actions taken
        '''

        loss = 0
        for i in range(len(prev_order_phase_labels)):
            phase = prev_order_phase_labels[i]
            if power in phase.keys() and phase[power] != [] and phase[power] != None:
                provinces = [order.split()[1] for order in phase[power]]

                # labels for province at specific phase
                order_indices = [ORDER_DICT[order] for order in phase[power]]
                one_hots = tf.one_hot(order_indices,depth=13042)

                # predictions for province at specific phase
                # print(probs.shape)
                # predictions = tf.gather_nd(probs,zip([i for j in range(len(provinces))],[position_lists.index(province) for province in provinces]))
                predictions = tf.convert_to_tensor([probs[i][position_lists.index(province)] for province in provinces], dtype=tf.float32)
                loss += tf.reduce_mean(categorical_crossentropy(one_hots, predictions))

        loss /= len(prev_order_phase_labels)
        return loss

    def train(self, state_inputs, prev_order_inputs, prev_orders_game_labels, season_names, board_dict_list):
        # Set up tracking metrics and logging directory
        train_loss = tf.keras.metrics.Mean('train_loss', dtype=tf.float32)
        current_time = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
        train_log_dir = 'data/logs/' + current_time + '/train'
        train_summary_writer = tf.summary.create_file_writer(train_log_dir)

        for i in range(len(state_inputs)):
            # Finds winning power to use in network
            last_turn = board_dict_list[i][-1]
            prov_num_dict = defaultdict(int)
            for province in last_turn:
                if "unit_power" in last_turn[province].keys():
                    prov_num_dict[last_turn[province]["unit_power"]] += 1
            power = max(prov_num_dict, key=prov_num_dict.get)

            # Parsing just the season(not year)
            # Not sure about these conversions
            powers_seasons = []
            curr_board_dict = board_dict_list[i]
            # extracting seasons and powers for film
            for j in range(len(season_names[i])):
                # print(season_names[i][j][0])
                powers_seasons.append(
                    SEASON[season_names[i][j][0]] + UNIT_POWER[power])
            # print(powers_seasons)

            # casting encoder and decoder inputs to floats
            powers_seasons = tf.convert_to_tensor(powers_seasons,
                                                  dtype=tf.float32)
            state_input = tf.convert_to_tensor(state_inputs[i],
                                               dtype=tf.float32)
            order_inputs = tf.convert_to_tensor(prev_order_inputs[i],
                                                dtype=tf.float32)
            season_input = season_names[i]

            with tf.GradientTape() as tape:
                try:
                    # applying SL model
                    orders_probs, position_lists = self.call(state_input,
                                                              order_inputs,
                                                              powers_seasons,
                                                              season_input,
                                                              curr_board_dict,
                                                              power)
                    # print(orders_probs.shape)
                    if orders_probs.shape[0] != 0:
                        orders_probs = tf.transpose(orders_probs, perm=[2, 0, 3, 1])
                        orders_probs = tf.squeeze(orders_probs)

                        # computing loss for probabilities
                        game_loss = self.loss(prev_orders_game_labels[i],
                                               orders_probs, position_lists, power)
                        print(game_loss)
                        # Add to loss tracking and record loss
                        train_loss(game_loss)
                        with train_summary_writer.as_default():
                            tf.summary.scalar('loss', train_loss.result(), step=i)
                        # optimizing
                        gradients = tape.gradient(game_loss,
                                                  self.trainable_variables)
                        self.optimizer.apply_gradients(
                            zip(gradients, self.trainable_variables))
                except:
                    continue
    # def get_orders(self, game, power_names):


if __name__ == "__main__":
    processor = process.Process("data/standard_no_press.jsonl")
    # setting weights of model
    new_weights_file = open("sl_weights.pickle","rb+")
    new_weights = pickle.load(new_weights_file)
    new_weights_file.close()
    sl_model = SL_model(16, 16)
    state_inputs, prev_order_inputs, prev_orders_game_labels, season_names, supply_center_owners, board_dict_list = processor.get_data(num_games=400)
    set_sl_weights(new_weights, sl_model, state_inputs, prev_order_inputs, prev_orders_game_labels, season_names, board_dict_list)
    # initializing model with 16 layers of each as in original paper
    for i in range(1000):
        state_inputs, prev_order_inputs, prev_orders_game_labels, season_names, supply_center_owners, board_dict_list = processor.get_data(num_games=100)
        sl_model.train(state_inputs, prev_order_inputs, prev_orders_game_labels, season_names, board_dict_list)

        # saving weights of SL model
        weights_file = open("sl_weights.pickle","wb+")
        pickle.dump(sl_model.get_weights(), weights_file)
        weights_file.close()
        print("Chunk %d" % (i))

    # setting weights of model
    new_weights_file = open("sl_weights.pickle","rb+")
    new_weights = pickle.load(new_weights_file)
    new_sl_model = SL_model(16, 16)
    set_sl_weights(new_weights, new_sl_model, state_inputs, prev_order_inputs, prev_orders_game_labels, season_names, board_dict_list)