Add support for parsing replays to gbx.py; sort tracks based on replays; fix important bug with data augumentation; new position model; apply weights to block model's output; add gui.py

Author: donadigo

README.md

@@ -18,12 +18,12 @@ This will generate a track using the provided block and position models that wil
 * Keras
 * python-lzo (through pip)
 * numpy
-* Not required: pygame or Gtk+3 and GLib for track visualization
+* Not required: Gtk+3 and GLib for track visualization
 
 ## Dataset
 This repo doesn't contain the dataset itself used to train the models in the `models/` directory as it is unusual to provide entire datasets with code in one repo. There is however a preprocessed version of the dataset used in the `data/train_data.pkl` file that you can use for futher training.
 
-The file contains roughly 3000 tech tracks downloaded directly from [TMX](https://tmnforever.tm-exchange.com/) and preprocessed such that they contain only the simplified version of tracks of each map. The maps themselves were downloaded using these filters: type: tech, order: awards (most), length: ~= 1m.
+The file contains roughly 3000 tech tracks downloaded directly from [TMX](https://tmnforever.tm-exchange.com/) and preprocessed such that they contain only the simplified version of tracks of each map. The maps themselves were downloaded using these filters: style: tech, order: awards (most), length: ~= 1m.
 
 ## Neural Network Architecture
 We can represent a track as a sequence of block placements. Each block consists of 3 main features:
@@ -44,8 +44,6 @@ Since we want to predict the next block in the sequence we ask the block model t
 The position model's output is two features: the vector to add to the position of last block to get a new position of the new block and the rotation of the new block.
 Their loss function is mean squared error and softmax respectively.
 
-![Visualization](/docs/TMTrackArch.png)
-
 ## Training
 It is recommended to have a dedicated GPU for training the nets, otherwise training process will be very slow.
 
@@ -62,4 +60,4 @@ python3 -i train_pos.py -g -l models/position_model_1024_512.h5
 Invoking either `train_blocks.py` or `train_pos.py` with the `-l` option will automatically 
 use model checkpointing to save new models with the model filename that was loaded.
 
-`livebuild.py` allows to dynamically generate tracks, it has a Gtk+3 UI to visualize how the track currently looks like. To fully evaluate model's performance, it's recommended to use `build.py` and see the tracks generated in the game itself.
+`livebuild.py` allows to dynamically generate tracks, it has a Gtk+3 UI to visualize how the track currently looks like. To fully evaluate model's performance, it's recommended to use `build.py` and see the tracks generated in the game itself.

blocks.py

@@ -361,7 +361,9 @@
 
     # 'StadiumRoadDirtHighToRoad': 358,
     # 'StadiumDirtHill': 359,
-    # 'StadiumPlatformBiSlope2StartSmall': 360
+    # 'StadiumPlatformBiSlope2StartSmall': 360,
+    # 'StadiumWaterClip': 361,
+    # 'StadiumDirtClip': 362
 }
 
 # TODO: needs a separate dict / list for TM2 Stadium
@@ -711,10 +713,12 @@ def is_multilap(name):
 
 BID, BX, BY, BZ, BROT, BFLAGS = range(6)
 EMPTY_BLOCK = (0, 0, 0, 0, 0)
-BASE_BLOCKS = list(range(6, 97+1)) + \
-    list(range(105, 186+1)) + list(range(196, 233+1))
-GROUND_BLOCKS = [20] + list(range(196, 233+1))
+BASE_BLOCKS = list(range(6, 98+1)) + \
+    list(range(105, 120+1)) + list(range(127, 224+1)) + list(range(227, 233+1))
+GROUND_BLOCKS = [20] + list(range(196, 233+1)) + \
+    list(range(253, 263+1)) + list(range(272, 277+1)) + list(range(287, 292+1))
 TRANSITION_BLOCKS = [130, 214, 215]
+ROAD_BLOCKS = [6, 198, 264, 265, 266, 270, 271]
 START_LINE_BLOCK = BLOCKS['StadiumRoadMainStartLine']
 FINISH_LINE_BLOCK = BLOCKS['StadiumRoadMainFinishLine']
 MULTILAP_LINE_BLOCK = BLOCKS['StadiumRoadMainStartFinishLine']

build.py

@@ -23,7 +23,7 @@
 from sklearn.preprocessing import MinMaxScaler
 
 from blocks import BID, BLOCKS, BROT, BX, BY, BZ
-from track_utils import fit_position_scaler
+from track_utils import fit_data_scaler
 from config import NET_CONFIG
 from builder import Builder
 from savegbx import save_gbx
@@ -40,7 +40,7 @@
 pattern_data_file = open(NET_CONFIG['patterns_fname'], 'rb')
 pattern_data = pickle.load(pattern_data_file)
 
-scaler = fit_position_scaler(train_data)
+scaler = fit_data_scaler(train_data)
 
 
 def progress_callback(completed, total):
@@ -56,7 +56,7 @@ def progress_callback(completed, total):
 pos_model = load_model(args.pos_model)
 
 builder = Builder(block_model, pos_model,
-                     NET_CONFIG['lookback'], train_data, pattern_data, scaler, temperature=args.temperature)
+                  NET_CONFIG['lookback'], train_data, pattern_data, scaler, temperature=args.temperature)
 
 track = builder.build(args.length, verbose=False, save=False,
                       progress_callback=progress_callback)

builder.py

@@ -4,8 +4,9 @@
 import numpy as np
 
 import blocks as bl
-from track_utils import intersects, occupied_track_positions, rotate_track_tuples, dist
+from track_utils import intersects, occupied_track_vectors, rotate_track_tuples, dist
 from blocks import BID, BROT, BX, BY, BZ
+from tech_block_weights import TECH_BLOCK_WEIGHTS
 
 POS_LEN = 3
 ROTATE_LEN = 4
@@ -22,15 +23,15 @@ def __init__(self, block_model, position_model, lookback, seed_data, pattern_dat
         self.inp_len = len(bl.BLOCKS) + POS_LEN + ROTATE_LEN
         self.temperature = temperature
         self.pattern_data = pattern_data
+        self.max_map_size = (32, 32, 32)
+        self.running = False
 
     @staticmethod
     def random_start_block():
         return (bl.START_LINE_BLOCK, 0, 0, 0, random.randrange(0, 4))
 
-    # Source:
-    # https://github.com/keras-team/keras/blob/master/examples/lstm_text_generation.py#L66
-    # Helper function to sample an index from a probability array
     def sample(self, preds):
+        # helper function to sample an index from a probability array
         preds = np.asarray(preds).astype('float64')
         preds = np.log(preds) / self.temperature
         exp_preds = np.exp(preds)
@@ -68,8 +69,7 @@ def decoded_track(track, start_pos=(0, 0, 0)):
                     block[BROT])
         return d
 
-    @staticmethod
-    def unpack_position_preds_vector(preds):
+    def unpack_position_preds_vector(self, preds):
         pos_vec = [int(round(axis)) for axis in preds[0][0]]
         pos_rot = np.argmax(preds[1][0])
         return pos_vec, pos_rot
@@ -79,6 +79,7 @@ def predict_next_block(self, X_block, X_position, block_override=-1, blacklist=[
             next_block = block_override
         else:
             block_preds = self.block_model.predict(X_block)[0]
+            block_preds = block_preds * TECH_BLOCK_WEIGHTS
             block_preds = np.delete(
                 block_preds, [bid - 1 for bid in blacklist])
 
@@ -136,7 +137,7 @@ def prepare_inputs(self):
         return (X_block, X_position)
 
     def position_track(self, track):
-        occ = occupied_track_positions(track)
+        occ = occupied_track_vectors(track)
         min_x = min(occ, key=lambda pos: pos[0])[0]
         min_y = min(occ, key=lambda pos: pos[1])[1] - 1
         min_z = min(occ, key=lambda pos: pos[2])[2]
@@ -145,6 +146,13 @@ def position_track(self, track):
         max_y = max(occ, key=lambda pos: pos[1])[1] - 1
         max_z = max(occ, key=lambda pos: pos[2])[2]
 
+        cx = 32 - (max_x - min_x + 1)
+        if cx > 0:
+            cx = random.randrange(0, cx)
+        cz = 32 - (max_z - min_z + 1)
+        if cz > 0:
+            cz = random.randrange(0, cz)
+
         min_x = 0 if min_x >= 0 else min_x
         min_y = 0 if min_y >= 0 else min_y
         min_z = 0 if min_z >= 0 else min_z
@@ -153,26 +161,44 @@ def position_track(self, track):
         max_y = 0 if max_y < 32 else max_y - 31
         max_z = 0 if max_z < 32 else max_z - 31
 
+        xoff = min_x - max_x
+        yoff = min_y - max_y
+        zoff = min_z - max_z
+
         p = []
         for block in track:
-            p.append((block[BID], block[BX] - min_x - max_x, block[BY] -
-                      min_y - max_y, block[BZ] - min_z - max_z, block[BROT]))
+            p.append((block[BID], block[BX] - xoff + cx, block[BY] -
+                      yoff, block[BZ] - zoff + cz, block[BROT]))
 
         return p
 
     def exceeds_map_size(self, track):
-        occ = occupied_track_positions(track)
+        occ = occupied_track_vectors(track)
         min_x = min(occ, key=lambda pos: pos[0])[0]
-        min_y = min(occ, key=lambda pos: pos[1])[1] - 1
+        min_y = min(occ, key=lambda pos: pos[1])[1]
         min_z = min(occ, key=lambda pos: pos[2])[2]
 
         max_x = max(occ, key=lambda pos: pos[0])[0]
-        max_y = max(occ, key=lambda pos: pos[1])[1] - 1
+        max_y = max(occ, key=lambda pos: pos[1])[1]
         max_z = max(occ, key=lambda pos: pos[2])[2]
 
-        return max_x - min_x > 32 or max_y - min_y > 32 or max_z - min_z > 32
+        return max_x - min_x + 1 > self.max_map_size[0] or max_y - min_y + 1 > self.max_map_size[1] or max_z - min_z + 1 > self.max_map_size[2]
+
+    def stop(self):
+        self.running = False
+
+    def get_y_locked(self):
+        for block in self.track:
+            if block[BID] in bl.GROUND_BLOCKS:
+                return True
+
+        return False
 
     def build(self, track_len, use_seed=False, failsafe=True, verbose=True, save=True, progress_callback=None):
+        self.running = True
+
+        # self.max_map_size = (random.randrange(
+        #     12, 32+1), random.randrange(5, 10), random.randrange(12, 32+1))
         if use_seed:
             self.track = self.sample_seed(3)
         else:
@@ -182,18 +208,21 @@ def build(self, track_len, use_seed=False, failsafe=True, verbose=True, save=Tru
         end = False
         current_min_y = 0
         while len(self.track) < track_len:
-            if len(blacklist) >= 5 or (len(blacklist) == 1 and end):
+            if not self.running:
+                return None
+
+            if len(blacklist) >= 10 or (len(blacklist) == 1 and end):
                 if verbose:
                     print('More than 10 fails, going back.')
 
                 if end:
-                    back = 10
+                    back = 5
                 else:
                     back = random.randrange(1, 4)
-                # Remove some last blocks
-                for _ in range(back):
-                    if len(self.track) > 1:
-                        del self.track[-1]
+
+                end_idx = min(len(self.track) - 1, back)
+                if end_idx > 0:
+                    del self.track[-end_idx:len(self.track)]
 
                 end = False
                 blacklist = []
@@ -210,35 +239,42 @@ def build(self, track_len, use_seed=False, failsafe=True, verbose=True, save=Tru
             decoded = self.decoded_track(
                 self.track + [next_block], start_pos=(0, 0, 0))
 
-            # Do not exceed map size
             if failsafe:
+                # Do not exceed map size
                 if self.exceeds_map_size(decoded):
                     blacklist.append(next_block[BID])
                     continue
 
                 if decoded[-1][BY] > current_min_y:
+                    # TODO: encode ground bit in the position network
                     if decoded[-1][BID] == 6 and decoded[-2][BID] == 6 and dist(decoded[-1][BX:BZ+1], decoded[-2][BX:BZ+1]) > 1:
                         blacklist.append(next_block[BID])
                         continue
 
+                    # Wants to put a ground block higher than ground
                     if next_block[BID] in bl.GROUND_BLOCKS:
-                        blacklist.append(next_block[BID])
+                        blacklist.extend(bl.GROUND_BLOCKS)
                         continue
 
                 if (intersects(decoded[:-1], decoded[-1]) or  # Overlaps the track
-                    (next_block[BID] in range(99, 104+1) or next_block[BID] in range(121, 126+1)) or
                         (next_block[BID] == bl.FINISH_LINE_BLOCK and not end)):  # Tries to put finish before desired track length
                     blacklist.append(next_block[BID])
                     continue
 
-                if self.score_prediction(self.track[-1], next_block) < 4:
+                if self.score_prediction(self.track[-1], next_block) < 5:
                     blacklist.append(next_block[BID])
                     continue
 
             blacklist = []
 
-            if decoded[-1][BY] < current_min_y:
-                current_min_y = decoded[-1][BY]
+            occ = occupied_track_vectors([decoded[-1]])
+            min_y_block = min(occ, key=lambda x: x[BY])[BY]
+            if min_y_block < current_min_y:
+                if self.get_y_locked():
+                    blacklist.append(next_block[BID])
+                    continue
+
+                current_min_y = min_y_block
 
             self.track.append(next_block)
             if len(self.track) >= track_len - 1:
@@ -252,5 +288,4 @@ def build(self, track_len, use_seed=False, failsafe=True, verbose=True, save=Tru
 
         result_track = self.position_track(
             self.decoded_track(self.track, (0, 0, 0)))
-        pickle.dump(result_track, open('generated-track.bin', 'wb+'))
         return result_track

bytereader.py

@@ -49,15 +49,27 @@ def get_bytes_file(self, num_bytes):
     def get_bytes_generic(self, num_bytes):
         return self.data[self.pos:self.pos + num_bytes]
 
+    def read_int32(self):
+        return self.read(4, 'i')
+
     def read_uint32(self):
         return self.read(4, 'I')
 
-    def read_int32(self):
-        return self.read(4, 'i')
+    def read_int16(self):
+        return self.read(2, 'h')
 
     def read_uint16(self):
         return self.read(2, 'H')
 
+    def read_int8(self):
+        return self.read(1, 'b')
+
+    def read_float(self):
+        return self.read(4, 'f')
+
+    def read_vec3(self):
+        return (self.read_float(), self.read_float(), self.read_float())
+
     def read_string(self):
         strlen = self.read_uint32()
         return self.read(strlen, str(strlen) + 's').decode('utf-8')
@@ -77,7 +89,7 @@ def read_info(self, info, f):
     def skip(self, num_bytes):
         self.pos += num_bytes
 
-    def read_string_loopback(self):
+    def read_string_lookback(self):
         if not self.seen_loopback:
             self.read_uint32()
 

config.py

@@ -1,6 +1,6 @@
 NET_CONFIG = {
     'batch_size': 128,
     'lookback': 20,
-    'train_fname': 'data/train_data.pkl',
-    'patterns_fname': 'data/pattern_data.pkl'
+    'train_fname': 'data/replay_train_data.pkl',
+    'patterns_fname': 'data/pattern_data_replay.pkl'
 }