Tile.Level/TileManor/scripts/rnn.py

#!/usr/local/bin/python3
# -*- coding: utf-8 -*-

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.data as tudata
import numpy as np
import sys
import json
import os

#将数据输出到一个json文件，便于在tiled进行观察，可以是单层或者多层的，多层的话，层级按照从小到大排列（层级是布局层级，不是视觉层级）
def output_to_json(slices, suffix="ex"):
    with open("/Users/xulianxin/Documents/develop/game/TileMatch/TileManor.Lv/TileManor/templates/tm_0000.json", 'r') as f:
        json_temp = json.load(f)
        for i in range(0, len(slices)):
            with open("/Users/xulianxin/Documents/develop/game/TileMatch/TileManor.Lv/TileManor/slices/slice_%d_%s.json" % (i, suffix), "w") as out_f:
                json_temp["layers"][0]["data"] = slices[i]
                json.dump(json_temp, out_f)

# 解析一个关卡文件，得到各层的数据，是按照视觉层次划分的
def get_each_view_slice(jsonLv):
    jsonData = ""
    tilesByPos = {}
    maxZ = 0
    # 处理各层的tile数据
    with open(jsonLv, 'r') as f:#, encoding='utf-8'
        jsonData = json.load(f)
        # 得到宽高信息
        w = int(jsonData['width'])
        h = int(jsonData['height'])
        for item in jsonData['layers']:
            name = item['name']
            if not name.startswith('Tile_'):
                # stacked
                continue
            z = int(name[5:])
            data = item['data']
            for i in range(0, len(data)):
                if data[i] == 0:
                    continue
                tile_x = (int)(i % w)
                tile_y = (int)(i / h)
                tile_z = z
                # 一些统计信息
                tilesByPos[(tile_x, tile_y, tile_z)] = [1,0]    # [tile类型，zView--视觉层级，默认是0]
                if tile_z > maxZ:
                    maxZ = tile_z
    # 根据上面的布局信息，计算每个tile的几个信息：
    # 1. 每个tile的视觉层级（不同于上面z的信息，那是一个布局信息，相同的z可能出在不同的视觉层级）
    # 从maxZ开始，逐层向下计算；对于每一个位置，如果该位置有tile，则计算该tile的视觉层级
    # 对于每一个tile，如果其上方（从该tile的z+1层，一直到maxZ层）有tile，则其视觉层级为上方tile的视觉层级+1
    slice_by_view_z = [[0]*w*h for i in range(maxZ+1)]
    for z in range(0, maxZ+1):
        z = maxZ - z
        for x in range(0, w+1):
            for y in range(0, h+1):
                if (x,y,z) in tilesByPos:
                    tile = tilesByPos[(x, y, z)]
                    adjs = [(x,y),(x,y+1),(x,y-1),(x+1,y),(x+1,y+1),(x+1,y-1),(x-1,y),(x-1,y-1),(x-1,y+1)]
                    # 确定视觉层级
                    zvMax = -1
                    for zup in range(z+1, maxZ+1):
                        for adj in adjs:
                            if (adj[0],adj[1],zup) in tilesByPos:
                                zv = tilesByPos[adj[0],adj[1],zup][1]
                                if zv > zvMax:
                                    zvMax = zv
                    zv = zvMax + 1
                    tilesByPos[(x, y, z)][1] = zv
                    # 记录到slice_by_view_z中
                    slice_by_view_z[zv][y*w+x] = tilesByPos[(x, y, z)][0]
    return slice_by_view_z

# 测试获取视觉层级的数据是否ok
def test_get_each_view_slice():
    jsonLv = "/Users/xulianxin/Documents/develop/game/TileMatch/TileManor.Lv/TileManor/templates/tm_0016.json"
    slice_by_view_z = get_each_view_slice(jsonLv)
    output_to_json(slice_by_view_z)

##########################################
# PyTorch RNN
class TileMatchRNN(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers, output_size):
        super(TileMatchRNN, self).__init__()
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.rnn = nn.GRU(input_size, hidden_size, num_layers, batch_first=True)
        self.fc = nn.Linear(hidden_size, output_size)
        # 检查是否有可用的GPU
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    def forward(self, x):
        # Add an extra dimension for the batch size if necessary
        if x.dim() == 2:
            x = x.unsqueeze(0)
        out, h = self.rnn(x)
        out = self.fc(out[:, -1, :])
        return out, h
    
# 全连接神经网络
class TileMatchFullyConnectedNetwork(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers, output_size):
        super(TileMatchFullyConnectedNetwork, self).__init__()
        self.layer1 = nn.Linear(input_size, hidden_size)
        self.layer2 = nn.Linear(hidden_size, output_size)
        # 检查是否有可用的GPU
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    def forward(self, x):
        x = torch.relu(self.layer1(x))
        x = self.layer2(x)
        return x
    
# 变分自编码器
class TileMatchVAE(nn.Module):
    def __init__(self, input_size, hidden_size, latent_size):
        super(TileMatchVAE, self).__init__()
        self.encoder = nn.Sequential(
            nn.Linear(input_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, latent_size * 2)  # We need mean and variance for each latent variable
        )
        self.decoder = nn.Sequential(
            nn.Linear(latent_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, input_size),
            nn.Sigmoid()  # To get outputs in the range [0, 1]
        )

    def reparameterize(self, mu, logvar):
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        return mu + eps * std

    def forward(self, x):
        h = self.encoder(x)
        mu, logvar = h.chunk(2, dim=1)
        z = self.reparameterize(mu, logvar)
        return self.decoder(z), mu, logvar

class TileMatchDataset(tudata.Dataset):
    def __init__(self, directories):
        self.data = []
        for directory in directories:
            for filename in os.listdir(directory):
                filepath = os.path.join(directory, filename)
                if os.path.isfile(filepath) and filepath.endswith('.json'):
                    slices = get_each_view_slice(filepath)
                    if len(slices) > 1:
                        for i in range(0, len(slices)-1):
                            self.data.append([slices[i], slices[i+1]])

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        input_sequence = self.data[idx][0]
        target_sequence = self.data[idx][1]
        return torch.tensor(input_sequence, dtype=torch.float32), torch.tensor(target_sequence, dtype=torch.float32)

# 模型的参数
input_size = 900
hidden_size = 50
num_layers = 2
output_size = 900

def train_NN():
    # Initialize model, loss function, optimizer
    # model = TileMatchRNN(input_size, hidden_size, num_layers, output_size)
    # criterion = nn.BCEWithLogitsLoss()  # 适合二分类问题
    # optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

    model = TileMatchVAE(input_size, 500, 20)
    criterion = nn.BCEWithLogitsLoss()  # 适合二分类问题
    optimizer = torch.optim.SGD(model.parameters(), lr=0.001)

    # 准备数据集
    batch_size = 1
    directories = ["/Users/xulianxin/Documents/develop/game/TileMatch/TileManor.Lv/TileManor/tf_templates",]
    dataset = TileMatchDataset(directories)
    train_loader = tudata.DataLoader(dataset, batch_size=batch_size, shuffle=True)

    num_epochs = 10
    for epoch in range(num_epochs):
        for data in train_loader:  # Assume train_loader provides (input, target) pairs
            inputs, targets = data
            # outputs, _ = model(inputs)
            x_recon, mu, logvar = model(inputs)
            recon_loss = F.binary_cross_entropy(x_recon, inputs, reduction='sum')
            kl_div = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
            loss = recon_loss + kl_div

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
    return model

if __name__ == '__main__':
    # test_get_each_view_slice()
    n = train_NN()
    torch.save(n, '/Users/xulianxin/Documents/develop/game/TileMatch/TileManor.Lv/TileManor/nn_VAE.model')
    # n = torch.load('/Users/xulianxin/Documents/develop/game/TileMatch/TileManor.Lv/TileManor/nn_VAE.model')
    slices = get_each_view_slice("/Users/xulianxin/Documents/develop/game/TileMatch/TileManor.Lv/TileManor/templates/tm_0015.json")
    n.eval()
    # 准备输入数据
    initial_input = torch.tensor(slices[0], dtype=torch.float32).unsqueeze(0)  # (1, 1, 900)

    # 初始化隐藏状态
    h0 = torch.zeros(num_layers, initial_input.size(0), hidden_size)

    # 使用模型进行预测 生成6层
    slices = []
    for i in range(1):
        with torch.no_grad():  # 在评估模式下，不需要计算梯度
            # output, hn = n(initial_input)
            output, mu, logvar = n(initial_input)
            initial_input = output
            # 应用阈值判断将浮点数转化为0或1
            threshold = 0.9
            output_layout = (torch.sigmoid(output) >= threshold).int()
            # 将 output_layout 转化成一个[]
            # slices.append([int(d) for d in output_layout[0]])
            slices.append([int(d) for d in output_layout[0]])
    output_to_json(slices, "nn_vae")