Deep Reinforcement Learning超簡單入門項目 Pytorch實現接水果遊戲AI

學習過傳統的監督和無監督學習方法後，我們現在已經可以自行開發機器學習系統來解決一些實際問題了。我們能實現一些事件的預測，一些模式的分類，還有數據的聚類等項目。但是這些好像和我們心目中的人工智能仍有差距，我們可能會認爲，人工智能是能理解人類語言，模仿人類行爲，並做到人類難以完成的工作的機器。所謂KNN、決策樹分類器，好像只是代替人類進行一些簡單的工作。
但今天，我們似乎在強化學習的領域找到了通往真正的人工智能的大門。經過強化學習訓練的AI，似乎已經可以做到人類做不到的事情。chat bot可以生成逼真的語言，GAN可以進行藝術創作，甚至有些AI可以在星際爭霸上打贏人類玩家。這些工作都與強化學習分不開關係。
我們之前實現的學習井字棋強化學習似乎已經能在完全沒有訓練資料的情況下學會下棋，甚至可以在不犯任何失誤的時候找到人類的破綻。雖然井字棋這個遊戲非常簡單，對20000不到的狀態空間，我們可以直接在蒙特卡洛樹中搜索出最好決策，但是這依然體現了強化學習的潛力。

強化學習與深度學習

在之前的學習過程中，我們學習了表格型的Q-Learning，在表格型的Q-Learning方法的學習過程中，我們逐漸會形成一張表格。在許多簡單的問題中，這種表格型的Q-Learning方法是比較實用的，但是當我們所處理的問題具有較大的狀態集合動作集時，這種表格型的方法就顯得十分的低效了。此時我們需要一種新的模型方法來處理這種問題，所以出現了結合了神經網絡的Q-Learning方法，Deep Q-Learning(DQN),通過在探索的過程中訓練網絡，最後所達到的目標就是將當前狀態輸入，得到的輸出就是對應它的動作值函數，也即$f(s)=Q(s,a)$ ,這個f就是訓練的網絡.

Deep Q learning

我們前面學習井字棋使用的學習方法是讓機器對弈，產生一條情節(episode)鏈，然後從後向前遍歷序列並更新Q值。每個狀態的更新公式都決定於該步action的獎賞和後續一個狀態的Q值。
$Q(x,a) = (1-\alpha)Q(x,a)+\alpha (R(x',a') + \gamma Q(x',a'))$
使用神經網絡來做deep的end to end Q學習時，這個問題就不是直接修改表，而是讓模型做一個迴歸。迴歸使用的誤差函數是MSE均方誤差。就是把上面的公式計算出來的新Q值當作迴歸目標，計算網絡輸出和它的均方誤差，然後用梯度方法更新一下就好了。

想讓DQL順利跑起來，還需要一些其他的工程技巧。
首先，我們在使用deep network擬合函數時，我們都是假定數據是獨立同分布的，並在一個有一定規模的數據集上運行梯度下降優化網絡。但是如果我們使用之前的方法，每進行一個對弈就用這個情節鏈去訓練一次。先不說想讓網絡適應這個情節需要幾次迭代更新，就算我們用很多次更新去讓網絡適應了這個經驗，神經網絡的特性也常常會出現在樣本數過少時的過擬合。而且下一次我們再用這個網絡去採樣時，就會出現完全不同分佈的採樣軌跡，這個問題的性質是無法保證收斂的。
爲此，2013年最早的DQL論文提出的方法是用兩個network，一個叫evaluate一個叫target。我們使用target網絡去對弈多次，採樣出多條探索軌跡形成數據集(這個探索次數大小自行調整)，再用這個數據集在evaluate網絡上訓練。這樣的數據集是獨立同分布的，網絡不容易過擬合，而且訓練震盪發散可能性降低，更加穩定。訓練得差不多之後，我們再把evaluate網絡的參數拷貝給target網絡，然後進入下一個epoch。

接水果

我們設計一個簡單的接水果遊戲來體現DQL與AI的工作邏輯，也許你覺得讓AI學習井字棋太簡單了，那麼現在我們真正讓AI來學習打電玩。
接水果遊戲的遊戲界面爲10x8，最上方每8個時間單位在隨機位置出現水果，每個時間單位後水果下降一格，每個時間單位AI或玩家可以控制最下方的3寬度的盤子向左或向右移動一格。
AI控制遊戲的方式是通過向前看一步，評估當前時刻進行任何操作獲得的預期獎勵，並做出最好決策。爲了避免AI在探索遊戲世界的過程中只探索當前模型給出的最優路徑，導致泛化性不足；我們也會設置一些比較小的噪聲，按照指數分佈去選擇優先級低一些的操作。

接水果遊戲實現

首先導入繪圖庫和numpy、torch等矩陣運算庫和深度學習庫。

import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.optim import Adam
from torch.utils.data import Dataset
import torch.utils.data as Data
import random
from collections import Counter

然後我們自己寫幾個函數實現接水果遊戲的邏輯。

def get_init_game():
    '''
    得到初始遊戲畫面
    '''
    game = torch.zeros((10,8))
    game[9,3:6] = 1.
    return (game,0)
    
def elapse(game,t,action):
    '''
    在game的遊戲畫面上推移一個時間單位，根據是action和t
    action可以是-1,0,1，代表向左，不動和向右
    如果t是8的倍數在遊戲畫面最上方產生新的水果
    水果下移一個像素，如果水果被接到就得到1 reward
    反之如果沒能接到就獲得-1 的reward
    界面上的0爲空地，1爲天上的水果或地上的盤子，
    2爲落在盤子裏的水果，3爲落在地上的水果
    '''
    game = game.clone()
    reward = 0.
    # 檢測地上有沒有爛水果
    game[9][torch.where(game[9]==2)] = 1
    game[9][torch.where(game[9]==3)] = 0
    # 根據action移動盤子
    if action==-1:
        game[9][0:7] = game[9][1:8].clone()
        game[9][7] = 0
    elif action==1:
        game[9][1:8] = game[9][0:7].clone()
        game[9][0] = 0
    # 判斷果子落地與否，落到盤子上
    if 1 in game[8]:
        fruit = torch.where(game[8]==1)
        if game[9][fruit] != 1:
            reward = -1.
            game[9][fruit] = 3
        else:
            reward = 1.
            game[9][fruit] = 2
        
    # 畫面下移一個像素
    game[1:9] = game[0:8].clone()
    game[0] = 0
    # 如果經過了8個時間，就隨機生成果子
    if t%8==0:
        #torch.seed = t
        idx = torch.randint(low = 0, high = 8, size = (1,))
        game[0][idx] = 1.
    t+=1
    return game,t,reward

第一個接口函數可以返回初始遊戲界面，第二個可是根據用戶或AI的action以及時間t，得到下一幀的時間推演函數。
然後我們定義讓AI探索遊戲和接管遊戲的控制函數

def successors(game,t):
    '''
    返回一個列表，列表元素是(game,t,reward)3元組
    '''
    ret = []
    # 不動
    ret.append(elapse(game,t,0))
    # 左
    if game[9][0]==0:
        ret.append(elapse(game,t,-1))
    # 右
    if game[9][-1]==0:
        ret.append(elapse(game,t,1))
    return ret

def poisson(n, lamda = 1):
    '''
    指數分佈，隨機生成0~n-1的整數，服從指數分佈
    '''
    probs = np.exp(-lamda*np.array(range(n+1)))
    probs /= probs.sum()
    for k in range(n):
        p = probs[k]/probs[k:].sum()
        if random.random()<p:
            return k
    return random.choice(range(n))

def policy(model, game, t, lamda):
    succ = successors(game, t)
    imgs = [x[0].unsqueeze(0).unsqueeze(0) for x in succ]
    imgs = torch.cat(imgs,axis = 0)
    values_pred = model(imgs.to("cuda")).view(-1)
    #print(values_pred)
    rank = torch.argsort(-values_pred)
    # 使用lamda係數的指數分佈，從最好到最壞隨機選擇下一步
    i = poisson(len(rank),lamda)
    #print(values_pred[i])
    return succ[rank[i]]

def episode(model, lamda):
    '''
    生成一整條情節鏈
    '''
    score = 0
    game,t = get_init_game()
    
    games = []
    rewards = []
    # 這裏預設情節鏈不超過一千幀
    for _ in range(1000):
        game,t,reward = policy(model, game, t, lamda)
        rewards.append(reward)
        games.append(game.unsqueeze(0).unsqueeze(0))
        if reward==-1:
            break
        score += reward
    return games,rewards,score

successors函數會告訴AI，如果它選擇某個動作則下一時刻會得到怎樣的局面。policy函數用model預測所有後繼的局面的優劣，並以指數分佈的概率從好到壞隨機選擇下一步(lambda係數越大，選擇最好的操作的信心就越高)
然後我們就可以用episode函數，讓AI完整地完一輪遊戲，直到它沒接到果子失敗，或者時間到達1000幀(如果AI已經學會了玩遊戲，他將會永遠玩下去，我們要讓他停下來)

模型

使用卷積神經網絡，爲了不丟失邊緣特徵，使用兩層3x3，1步長，1填充的卷積層。後接池化層和全連接層用於迴歸。事實上如果我們神經網絡的輸入可以不是一張圖片，也可以是連續的n個時間單位的n張圖片，這樣的設計允許網絡學習到時序的信息(比如在自動駕駛中，AI通過接收連續的幾個幀就能分析出當前車輛的運行速度)，但我們這個遊戲是很符合MDP假設的遊戲，一般來說只看一幀就可以做出決策，所以這裏並沒有使用這種輸入。

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        # 輸入:nx1x10x8張量
        self.conv1 = nn.Conv2d(1, 8, 3, padding=1)
        self.conv2 = nn.Conv2d(8, 16, 3, padding=1)
        # nx16x10x8張量
        self.maxpool = nn.MaxPool2d(2,2)
        # nx16x5x4張量
        self.fc = nn.Sequential(
            nn.Linear(16*5*4, 32),
            nn.ReLU(),
            nn.Linear(32, 1),
        )
        

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        x = self.maxpool(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x

優化方法與採樣

模型訓練使用隨機梯度下降，損失函數使用均方誤差。均方誤差很容易理解，因爲我們的問題本質上是迴歸問題，MSE是很好的迴歸用損失函數。使用SGD而不用adam等更先進的優化器的原因是，adam等優化器使用動量，它的使用強烈依賴於獨立同分布假設。而我們使用agent玩遊戲得到的採樣來訓練模型，這個採樣又和target model有絕對的相關性，是會隨時變動的。如果我們使用動量則很有可能把之前的錯誤優化方向擴大化，讓訓練不收斂。
如何採樣事實上也是有講究的，我們當然可以就讓target model不斷嘗試玩遊戲，並且把所有的錄像帶都用於訓練。但是有一點很重要，就是樣本的平衡與否。我們最開始的嘗試，會得到一些完全接不到水果的錄像，也會得到一些偶然接到水果的錄像，然後我們會用這些錄像帶訓練，逐漸的AI變得不那麼容易完全接不到。這時我們認爲"完全接不到"這個事件已經訓練到收斂了，因此就沒有用完全接不到的錄像帶再去train模型的意義。如果我們仍然保留這些完全接不到的錄像帶，則會讓模型得到錯誤的收斂方向。我們應該儘可能使用走得更遠的，分數更高的錄像帶來訓練模型。所以我們雖然會進行採樣，但是這個採樣需要有選擇性，我們應該重視表現得更好的錄像，而適當拋棄表現一般的錄像。
如果你沒有GPU，就把下面的cuda換成cpu。

LR = 0.002

eval_model = Net().to("cuda")
target_model = Net().to("cuda")

loss_func = nn.MSELoss()
optimizer = torch.optim.SGD(eval_model.parameters(), lr=LR)

訓練，我們會讓target模型採樣，得到一個錄像庫。每次訓練我們會從庫中隨機採樣並訓練eval模型。同時我們會定期把eval模型的參數拷貝給target模型，並定期讓target模型再做新的採樣，用新的episode擠掉舊的episode。

replays = []
# 採樣，保證樣本平衡
while(len(replays)<50):
    tpl = episode(target_model, 1)
    if tpl[2]==0:
        replays.append(tpl)

while(len(replays)<100):
    tpl = episode(target_model, 1)
    if tpl[2]!=0:
        replays.append(tpl)
        
random.shuffle(replays)

def train(discount = 0.9, num_iters = 10):
    # 產生新的情節鏈
    games,rewards,score = replays.pop(0)
    avg_score = sum(tpl[2] for tpl in replays)/100
    lamdas = [2**(i) for i in torch.linspace(3,0,11)]
    for lamda in lamdas:
        # 儘可能讓本次採樣的情節鏈得到的score>=當前score的平均值
        # 保證訓練樣本是儘可能優秀的樣本，才能讓模型不在訓練中退化
        # 試想，如果模型一直接收它的失敗經歷，則訓練是低效的
        tpl = episode(target_model, lamda)
        if tpl[2]>=avg_score:
            replays.append(tpl)
            break
    while(len(replays)<100):
        tpl = episode(target_model, 2)
        replays.append(tpl)
    
    running_loss = 0.
    running_score = 0.
    for step in range(num_iters):
        # 隨機選擇一個錄像進行訓練
        i = random.randint(0,49)
        games,rewards,score = replays[i]
        running_score += score
        rewards = torch.tensor(rewards).view(-1,1).cuda()
        games = torch.cat(games,axis = 0).cuda()
        Q_target = torch.zeros_like(rewards).cuda()
        
        # 計算Q_target(s) = gamme*Q(s')+reward
        with torch.no_grad():
            Q_target[:-1] = discount*target_model(games[1:])
        Q_target += rewards
        optimizer.zero_grad()
        output = eval_model(games)
        loss = loss_func(output,Q_target)
        running_loss += loss.item()
        loss.backward()
        optimizer.step()
    
    target_model.load_state_dict(eval_model.state_dict())
    return score,running_loss/num_iters

訓練

我們使用上面的方法，訓練10k個epoch，每個epoch內我們使用target model玩很多次遊戲，並選擇一個表現較好的錄像帶加入records隊列。訓練時，我們從隊列中隨機選擇一個錄像帶訓練，重複10次。
訓練中觀察到，損失值會隨着訓練的次數增加而減小，說明訓練是收斂的。隨着訓練次數增加，模型在遊戲上取得的分數也會增加，雖然有波動，但總體呈上升趨勢。當然，因爲玩每局遊戲都能走得更遠，越往後訓練單個epoch需要的時間也越長。

EPOCH = 10000
avg_score = 0
avg_loss = 0
step = 200

score_list = []
loss_list = []

for epoch in range(EPOCH):
    score,loss = train(discount = 0.85, num_iters = 10)
    avg_score+=score
    avg_loss+=loss
    if (epoch+1)%step==0:
        avg_loss/=step
        avg_score/=step
        score_list.append(float(avg_score))
        loss_list.append(float(avg_loss))
        print("Epoch %5d, loss: %6.3f, score: %.2f."%(epoch+1,avg_loss,avg_score))
        avg_loss=0
        avg_score=0
    # 因爲我們預設了episode不會超過1k幀，即最大的score爲124
    # 如果avg_score已經達到124很久則可以認爲收斂
    if len(score_list)>5 and score_list[-5]>=124:
        break

plt.figure()
plt.title('AI score')
plt.plot(score_list)
plt.xlabel('Epoch/(200)')
plt.ylabel('score')


plt.figure()
plt.title('Train loss')
plt.plot(loss_list)
plt.xlabel('Epoch/(200)')
plt.ylabel('MSEloss')

Epoch   200, loss:  0.105, score: 2.06.
Epoch   400, loss:  0.122, score: 3.53.
Epoch   600, loss:  0.122, score: 4.16.
Epoch   800, loss:  0.122, score: 4.34.
Epoch  1000, loss:  0.119, score: 5.16.
Epoch  1200, loss:  0.110, score: 8.65.
Epoch  1400, loss:  0.115, score: 12.30.
Epoch  1600, loss:  0.103, score: 10.51.
Epoch  1800, loss:  0.090, score: 9.11.
Epoch  2000, loss:  0.071, score: 16.80.
Epoch  2200, loss:  0.073, score: 12.28.
Epoch  2400, loss:  0.051, score: 14.52.
Epoch  2600, loss:  0.055, score: 16.41.
Epoch  2800, loss:  0.051, score: 16.34.
Epoch  3000, loss:  0.043, score: 27.86.
Epoch  3200, loss:  0.054, score: 21.86.
Epoch  3400, loss:  0.051, score: 10.62.
Epoch  3600, loss:  0.043, score: 18.92.
Epoch  3800, loss:  0.040, score: 13.55.
Epoch  4000, loss:  0.035, score: 13.74.
Epoch  4200, loss:  0.040, score: 14.37.
Epoch  4400, loss:  0.028, score: 22.42.
Epoch  4600, loss:  0.029, score: 55.03.
Epoch  4800, loss:  0.019, score: 70.56.
Epoch  5000, loss:  0.025, score: 46.06.
Epoch  5200, loss:  0.026, score: 38.34.
Epoch  5400, loss:  0.024, score: 59.89.
Epoch  5600, loss:  0.023, score: 81.97.
Epoch  5800, loss:  0.024, score: 66.81.
Epoch  6000, loss:  0.021, score: 77.39.
Epoch  6200, loss:  0.016, score: 88.88.
Epoch  6400, loss:  0.019, score: 84.44.
Epoch  6600, loss:  0.020, score: 71.55.
Epoch  6800, loss:  0.017, score: 79.44.
Epoch  7000, loss:  0.021, score: 71.22.
Epoch  7200, loss:  0.003, score: 124.00.
Epoch  7400, loss:  0.003, score: 124.00.
Epoch  7600, loss:  0.002, score: 124.00.
Epoch  7800, loss:  0.002, score: 124.00.
Epoch  8000, loss:  0.002, score: 124.00.

訓練到收斂後，執行下面的代碼就可以看到AI成功地在玩接水果遊戲玩了一千幀。

import time
from matplotlib.colors import ListedColormap
 
cmap_light = ListedColormap(['white', 'black', 'yellow', 'red'])
from IPython import display

games,rewards,score = episode(target_model, 5)
print("Score:",score)


for i in range(len(games)):
    plt.figure()
    img = games[i]
    img = img.squeeze()
    img = img.numpy()
    plt.imshow(img, cmap=cmap_light)
    plt.show()
    display.clear_output(wait=True)

訓練好的模型也可以保存起來

torch.save(target_model.state_dict(),'fruit-ai.cpk')
target_model.load_state_dict(torch.load('fruit-ai.cpk'))

總結

這個微型項目的初衷是簡單的入門Deep Reinforcement learning，因爲網絡上的強化學習入門項目無不是對硬件要求很高，需要訓練個幾百萬幀才能收斂的那種遊戲。更不用說每一幀都是100x100以上的大小，沒有硬件支持可能很難快速看到結果。這裏設計了超輕量級的遊戲，遊戲邏輯簡單，每一幀也只是10x8的大小。模型也是超輕量的，保存後可以看到，只用了46.5k的空間。
儘管如此，cpu上要把這個遊戲玩到收斂也需要接近五六萬次的採樣，約200k幀的訓練，花費2個小時以上。所以，我們還是該認識到，在深度學習時代算力是多麼重要的資源。