ResNet ：构建 AI 游戏机器人玩贪吃蛇

发布日期：2024-12-02 17:59:32 浏览次数： 2960

作者：AI零壹白洞

微信搜一搜，关注“AI零壹白洞”

小编分享过很多关于 AI agent 科普了，这篇讲从代码技术中给大家讲下，如何使用 PyTorch 通过模仿学习和运动分析对 Google Snake 进行 AI 构建 Agent，让它自己玩起来。

以往分享的原理和技术：

本文重点：

重点代码实现，从 0 构建 Google Snake AI agent，中间 python 库不是官方自带库，需要各位 pip install，或者用 ide 集成。
涉及 AI 知识科普。
代码已在 github，附录中。

本案例将使用模仿学习技术，即代理从人类对手那里学习如何做出决策。这意味着操作者首先要玩一段时间的游戏并收集数据。然后，人工智能使用这些数据来训练自己如何玩游戏并从人类那里获取策略。这就是所谓的模仿学习（不要与迁移学习混淆，后者核心思想是利用在一个任务上学习到的知识来提高在另一个相关任务上的性能）。

从技术角度来说，大家将训练一个 3D 卷积 ResNet 模型。这是因为我们想要捕捉运动以了解蛇前进的方向。为此，将一次为模型提供 4 帧游戏画面来推断运动。可以尝试只使用一帧并使用标准卷积模型（如 EfficientNet），但如果没有运动信息，效果就不会那么好。

模仿学习

模仿学习（Imitation Learning），也称为学习从演（Learning from Demonstration，LfD）或行为克隆（Behavioral Cloning，BC），是一种机器学习方法，它允许机器通过观察和模仿专家的行为来学习任务。模仿学习的一个关键优势是它不需要显式的奖励函数，这在许多复杂任务中是难以定义的。然而，它也有一些局限性，比如可能会学习到专家的非最优行为，或者在未见过的情况下表现不佳。

流程 & 框架

数据收集、准备数据、加载模型代码、训练模型、推理运行

数据收集

数据收集有很多种方法，这里使用的方法依赖于 selenium，这是一种自动化浏览器导航和控制的 Python 工具，selenium 能将屏幕截图保存为具有适当标签的图像的代码，主要是这个方便。

import base64import ioimport cv2from PIL import Imageimport numpy as npimport keyboardimport osfrom datetime import datetimefrom selenium import webdriverfrom selenium.webdriver.common.by import By# 初始化环境isExist = os.path.exists("captures")if isExist:    dir = "captures"    for f in os.listdir(dir):        os.remove(os.path.join(dir, f))else:    os.mkdir("captures")
current_key = "1"buffer = []# 收集用户键盘反馈def keyboardCallBack(key: keyboard.KeyboardEvent):    global current_key    if key.event_type == "down" and key.name not in buffer:        buffer.append(key.name)    if key.event_type == "up":        buffer.remove(key.name)    buffer.sort()    current_key = " ".join(buffer)    keyboard.hook(callback=keyboardCallBack)# 获取浏览器上下文driver = webdriver.Firefox()# 导航到 Google Snake gamedriver.get("<https://www.google.com/fbx?fbx=snake_arcade>")frame_stack = deque(maxlen=4)while True:    # 获取画布元素    canvas = driver.find_element(By.CSS_SELECTOR, "canvas")    # 获取画布数据，这里就是一张图，一帧    canvas_base64 = driver.execute_script(        "return arguments[0].toDataURL('image/png').substring(21);", canvas)    # Decode the base64 data to get the PNG image    canvas_png = base64.b64decode(canvas_base64)
    image = cv2.cvtColor(        np.array(Image.open(io.BytesIO(canvas_png))), cv2.COLOR_BGR2RGB)
    # 保存有用户键盘策略的图片和没有策略图片    if len(buffer) != 0:        cv2.imwrite(            "captures/" + str(datetime.now()).replace("-", "_").replace(":", "_").replace(" ", "_") + " "            + current_key + ".png", image, )    else:        cv2.imwrite(            "captures/" + str(datetime.now()).replace("-", "_").replace(":", "_").replace(" ", "_") + " n"            + ".png", image, )          # 计算每个标签预测值，用加权平均来计算            frame_stack.append(transformer(image))     input = torch.stack([*frame_stack], dim=1).to(device).squeeze().unsqueeze(0)
     if len(frame_stack) == 4:         with torch.inference_mode():         outputs = model(input).to(device)         preds = torch.softmax(outputs, dim=1).argmax(dim=1)
         if preds.item() != 0:           keyboard.press_and_release(label_keys[preds.item()])

运行此脚本后，会打开一个窗口来运行，然后即可开始玩游戏。在后台，脚本会不断保存游戏屏幕截图，并使用唯一时间戳和当前按下的键来命名图像。当没有按下任何键时，它会被标记为 n。

准备数据

将这些图像转换为带有文件名和相应操作的 csv 文件。

import pandas as pdimport matplotlib.pyplot as pltimport osimport csv import os
# 创建目录，保存包含标签和图像文件名的 CSV 文件。labels = []dir = 'captures'file_path = "data/labels_snake.csv"if not os.path.exists(file_path):    os.mkdir('data')# 读取文件名，从每张图片的文件名中提取按下的键。按下的键可以是左、右、上、下，或者根本没有键。    for f in os.listdir(dir):    key = f.rsplit('.',1)[0].rsplit(" ",1)[1]# 根据所按的键，将每幅图像分为四类：0 表示未按任何键，1 表示左，2 表示上，3 表示右，4 表示下。       if key=="n":        labels.append({'file_name': f, 'class': 0})    elif key=="left":        labels.append({'file_name': f, 'class': 1})    elif key=="up":        labels.append({'file_name': f, 'class': 2})    elif key=="right":        labels.append({'file_name': f, 'class': 3})    elif key=="down":        labels.append({'file_name': f, 'class': 4})# 创建标签文件，包含数据集来训练机器学习模型    field_names= ['file_name', 'class']with open('data/labels_snake.csv', 'w') as csvfile:    writer = csv.DictWriter(csvfile, fieldnames=field_names)    writer.writeheader()    writer.writerows(labels)

在此过程中，实际上是在创建一个带有相应标签的图像数据集，其中每个标签代表按下按键的方向。

加载模型代码

第一步先把数据导入 tensor，形式变成神经网络学习所要的数据形式。

from torch.utils.data import Dataset, DataLoader, WeightedRandomSamplerimport osfrom PIL import Imageimport torchfrom sklearn.model_selection import train_test_splitimport pandas as pdfrom torchvision.transforms import transforms, Compose, ToTensor, Resize, Normalize, CenterCrop, Grayscalefrom torch import nnfrom tqdm import tqdmfrom torchinfo import summaryimport numpy as npimport mathfrom torchvision.models.video import r3d_18, R3D_18_Weights, mc3_18, MC3_18_Weights

# 数据集由按时间顺序排列的四张图像堆栈组成。# 从数据集中提取的每个项目代表四帧序列，其中最后一帧与按键相关联。# 本质上，此数据集通过最后四帧捕获运动并将其与按键相关联。# 其中 stack_size 会影响后续权重，这个解释下，是堆叠在一起作为单个数据点的图像数量。class SnakeDataSet(Dataset):    # 包含有关图像和标签信息的数据集    def __init__(self, dataframe, root_dir, stack_size, transform=None):        self.stack_size = stack_size        self.key_frame = dataframe        self.root_dir = root_dir        self.transform = transform
    # 返回数据集的长度，即数据点的总数。    # 长度计算为 的长度key_frame减去 的三倍 stack_size。    # 这表明数据集预计包含图像序列，并且每个数据点由一堆图像组成。    def __len__(self):        return len(self.key_frame) - self.stack_size * 3
    # 获取索引idx并返回相应的数据点    def __getitem__(self, idx):        if torch.is_tensor(idx):            idx = idx.to_list()        try:            img_names = [os.path.join(self.root_dir, self.key_frame.iloc[idx + i, 0]) for i in range(self.stack_size)]            images = [Image.open(img_name) for img_name in img_names]            # 使用 tensor 提取人工数据标签            label = torch.tensor(self.key_frame.iloc[idx + self.stack_size, 1])            # 图片转化，将变换应用于序列中的每个图像。            if self.transform:                images = [self.transform(image) for image in images]        except:            img_names = [os.path.join(self.root_dir, self.key_frame.iloc[0 + i, 0]) for i in range(self.stack_size)]            images = [Image.open(img_name) for img_name in img_names]            # 如果遇到错误，使用数据第一个标签，处理兼容形式，也就是 left 形式            label = torch.tensor(self.key_frame.iloc[0 + self.stack_size, 1])            # 图片转化，将变换应用于序列中的每个图像。            if self.transform:                images = [self.transform(image) for image in images]        # 将图像沿着新维度堆叠 torch.stack(images, dim=1)，然后使用 删除单例维度squeeze()。这会产生一个表示堆叠图像的张量。        # 堆叠图像张量和标签张量作为元组返回。        return torch.stack(images, dim=1).squeeze(), label

到这里可能存在一个问题，大家都能猜到数据集这个平衡性。这是因为大多数时候，在玩游戏时，用户没有按任何键。因此，大头都是 0

。所以这里有 2 种策略去做后续改造，第一种就是要做个模拟随机性，随机数据样本，防止模型学习过程过拟合，走得后面全是 0。另一种就是调整权重，把获取分数降低来刺激模型选择其它。

本案例使用第二种

，RandomWeightedSampler 来解决问题。它将每个样本的权重作为输入，然后调整权重。

STACK_SIZE = 4BATCH_SIZE = 32
# 区分数据集和结果验证集，这里设置结果验证，占全样本 20%train, test = train_test_split(pd.read_csv("data/labels_snake.csv"), test_size=0.2, shuffle=False)classes = ["n", "left", "up", "right", "down"]
# 设置实验数据# 通过将所有计数的总和除以该类别出现的次数来计算每个类别的权重。labels_unique, counts = np.unique(train["class"], return_counts=True)class_weights = [sum(counts)/c for c in counts]# 下一步是通过为示例分配类权重来获取每个示例的权重。# 这是通过遍历数据集并根据其类标签为每个示例分配权重来完成的。example_weights = np.array([class_weights[l] for l in train['class']])# 根据堆栈大小滚动示例权重，因为与特定图像相关联的标签实际上是该图像索引的标签 + STACK_SIZE。# 这可确保根据其类标签为每个样本赋予正确的权重。# 0 建设数据 14w 数据集，所以这里需要减去 14w，要么权重其实比其它低的example_weights = np.roll(example_weights, -STACK_SIZE)sampler = WeightedRandomSampler(example_weights, len(train))
# 重新设置一遍测试验证结果集labels_unique, counts = np.unique(test["class"], return_counts=True)class_weights = [sum(counts)/c for c in counts]test_example_weights = np.array([class_weights[l] for l in test['class']])test_example_weights = np.roll(test_example_weights, -STACK_SIZE)test_sampler = WeightedRandomSampler(test_example_weights, len(test))
# 初始化模型 loaderdataset = SnakeDataSet(root_dir="captures", dataframe = train, stack_size=STACK_SIZE, transform=transformer)dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, sampler=sampler, drop_last= True)test_dataset = SnakeDataSet(root_dir="captures", dataframe = test, stack_size=STACK_SIZE,  transform=transformer)test_dataloader = DataLoader(test_dataset, batch_size=BATCH_SIZE, sampler = test_sampler, drop_last=True)

通过＋-加权法，计算数据集中每个样本的权重是机器学习任务中的关键步骤。它有助于平衡数据集并确保每个类别对学习过程的贡献相同。该过程包括将数据集拆分为训练集和测试集，获取唯一标签和每个标签的计数，计算每个类别的权重，并根据其类别标签为每个样本分配权重。

from torchvision.transforms import transforms, Compose, Normalize, CenterCropfrom torchvision.models.video import r3d_18, R3D_18_Weights, mc3_18, MC3_18_Weights
# 计算数据集图像的平均值和标准差def compute_mean_std(dataloader):    # source: <https://github.com/aladdinpersson/Machine-Learning-Collection/blob/master/ML/Pytorch/Basics/pytorch_std_mean.py>    # var[X] = E[X**2] - E[X]**2    channels_sum, channels_sqrd_sum, num_batches = 0, 0, 0    for batch_images, labels in tqdm(dataloader):  # (B,H,W,C)            batch_images = batch_images.permute(0,3,4,2,1)            channels_sum += torch.mean(batch_images, dim=[0, 1, 2, 3])            channels_sqrd_sum += torch.mean(batch_images ** 2, dim=[0, 1, 2,3])            num_batches += 1        mean = channels_sum / num_batches        std = (channels_sqrd_sum / num_batches - mean ** 2) ** 0.5        return mean, std
compute_mean_std(dataloader)# 将图像大小调整为 84x84，转换为张量并对其进行规范化。transformer = Compose([    Resize((84,84), antialias=True),    CenterCrop(84),    ToTensor(),    Normalize(mean =[ -0.7138, -2.9883,  1.5832], std =[0.2253, 0.2192, 0.2149])])         
# 使用 PyTorch 提供的 r3d 模型（ResNet架构）model = r3d_18(weights = R3D_18_Weights.DEFAULT)model.fc = nn.Linear(in_features=512, out_features=5, bias=True)summary(model, (32,3,4,84,84))

下面展示了r3d模型卷积过程，其中 [32, 64, 4, 42, 42] 这些就是卷积上下采集过程，总过 7 层。

训练模型

在这里，设置 10e-5 的学习率和 0.1 的权重衰减。

# 设置 epochs 次数，epochs 理解是一个批量处理，10w 次num_epochs = 2# 设置环境，损失函数交叉熵device = torch.device("cuda" if torch.cuda.is_available() else "cpu")optimizer = torch.optim.AdamW(model.parameters(), 10e-5, weight_decay=0.1)model.to(device)criterion = nn.CrossEntropyLoss()
for epoch in range(num_epochs):    total_loss = 0.0    correct_predictions = 0    total_samples = 0
    val_loss = 0.0    val_correct_predictions = 0    val_total_samples = 0
    # 开始训练    model.train()
    # 显示进度条    pbar = tqdm(dataloader, desc=f'Epoch {epoch + 1}/{num_epochs}', leave=True)    # 从 pbar 批处理种提取 inputs 和 labels    for inputs, labels in pbar:        inputs, labels = inputs.to(device), labels.to(device)
        outputs = model(inputs.to(device))        loss = criterion(outputs, labels)
        # 向后传播和优化        optimizer.zero_grad()        loss.backward()        optimizer.step()
        # 更新梯度和参数        total_loss += loss.item()        _, predicted = torch.max(torch.softmax(outputs,1), 1)        correct_predictions += (predicted == labels).sum().item()        total_samples += labels.size(0)
        # 更新损失率和精度        pbar.set_postfix({'Loss': total_loss / total_samples, 'Accuracy': correct_predictions / total_samples})        steps = steps + 1
    # 结果评估    model.eval()    with torch.inference_mode():        for inputs, labels in test_dataloader:            inputs, labels = inputs.to(device), labels.to(device)
            outputs = model(inputs.to(device))            loss = criterion(outputs, labels)
            # 更新损失率和精度            val_loss += loss.item()            _, predicted = torch.max(torch.softmax(outputs,1), 1)            val_correct_predictions += (predicted == labels).sum().item()            val_total_samples += labels.size(0)
    # 最终结果评估和输出    epoch_loss = val_loss / val_total_samples    epoch_accuracy = val_correct_predictions / val_total_samples    print(f'Epoch {epoch + 1}/{num_epochs}, Val Loss: {epoch_loss:.4f}, Val Accuracy: {epoch_accuracy:.4f}')    torch.save(model.state_dict(), "model_r3d.pth")