【深度強化學習】9. Policy Gradient實現中核心部分torch.distributions

【導語】:在深度強化學習第四篇中,講了Policy Gradient的理論。通過最終推導得到的公式,本文用PyTorch簡單實現以下,並且儘可能搞清楚torch.distribution的使用方法。代碼參考了LeeDeepRl-Notes中的實現。

1. 複習

\[\theta \leftarrow \theta+\eta \nabla \bar{R_\theta} \\\nabla \bar{R_\theta}=\frac{1}{N}\sum^N_{n=1}\sum^{T_n}_{t=1}R(\tau^n)\nabla log p_\theta(a_t^n|s_t^n) \]

\(\theta\)代表模型的參數,第一行公式代表了模型進行更新的方法,\(\eta\) 代表的是學習率。

第二行是推導得到的,和CrossEntropy可以對照着理解記憶。

2. Torch.Distributions

distributions包主要是實現了參數化的概率分佈和採樣函數。參數化是爲了讓模型能夠具有反向傳播的能力,這樣纔可以用隨機梯度下降的方法來進行優化。隨機採樣的話沒辦法直接反向傳播,有兩個方法,REINFORCE和pathwise derivative estimator。

Torch中提供兩個方法,sample()和log_prob(),就可以實現REINFORCE

\[\Delta \theta=\alpha r \frac{\partial \log p\left(a \mid \pi^{\theta}(s)\right)}{\partial \theta} \]

\(\theta\)是模型參數,\(\alpha\)代表的是學習率,r代表reward,\(p\left(a \mid \pi^{\theta}(s)\right)\)代表在狀態s下,使用策略\(\pi^{\theta}\)採取a動作的概率。

2.1 REINFORCE

實現的時候,會先從網絡輸出構造一個分佈,然後從分佈中採樣一個action,將action作用於環境,然後使用log_prob()函數來構建一個損失函數,代碼如下(PyTorch官方提供):

probs = policy_network(state)
# Note that this is equivalent to what used to be called multinomial
m = Categorical(probs)
action = m.sample()
next_state, reward = env.step(action)
loss = -m.log_prob(action) * reward
loss.backward()

對照一下,這個-m.log_prob(action)應該對應上述公式:\(\log p\left(a \mid \pi^{\theta}(s)\right)\), 加負號的原因是,在公式中應該是實現的梯度上升算法,而loss一般使用隨機梯度下降的,所以加個負號保持一致性。

2.2 PathWise Derivative Estimator

這是一種重參數化技巧,主要是通過調用rsample()函數來實現的,參數化隨機變量可以通過無參數隨機變量的參數化確定性函數來構造。參數化以後,採樣過程就變得可微分了,也就支持了網絡的後向傳播。實現如下(PyTorch官方實現):

params = policy_network(state)
m = Normal(*params)
# Any distribution with .has_rsample == True could work based on the application
action = m.rsample()
next_state, reward = env.step(action)  # Assuming that reward is differentiable
loss = -reward
loss.backward()

這樣的話,可以直接對-reward使用隨機梯度下降,因爲rsample後可微分,可以後向傳播。

3. 源碼

主要看agent對象的實現:

class PolicyGradient:
    def __init__(self, state_dim, device='cpu', gamma=0.99, lr=0.01, batch_size=5):
        self.gamma = gamma
        self.policy_net = FCN(state_dim)
        self.optimizer = torch.optim.RMSprop(
            self.policy_net.parameters(), lr=lr)
        self.batch_size = batch_size

    def choose_action(self, state):
        state = torch.from_numpy(state).float()
        state = Variable(state)
        probs = self.policy_net(state)
        m = Bernoulli(probs)
        action = m.sample()

        action = action.data.numpy().astype(int)[0]  # 轉爲標量
        return action

    def update(self, reward_pool, state_pool, action_pool):
        # Discount reward
        running_add = 0 # 就是那個有discount的公式

        for i in reversed(range(len(reward_pool))): # 倒數
            if reward_pool[i] == 0:
                running_add = 0
            else:
                running_add = running_add * self.gamma + reward_pool[i]
                reward_pool[i] = running_add
        # 得到G

        # Normalize reward
        reward_mean = np.mean(reward_pool)
        reward_std = np.std(reward_pool)
        for i in range(len(reward_pool)):
            reward_pool[i] = (reward_pool[i] - reward_mean) / reward_std
        # 歸一化

        # Gradient Desent
        self.optimizer.zero_grad()

        for i in range(len(reward_pool)): # 從前往後
            state = state_pool[i] 
            action = Variable(torch.FloatTensor([action_pool[i]]))
            reward = reward_pool[i]

            state = Variable(torch.from_numpy(state).float())
            probs = self.policy_net(state)
            m = Bernoulli(probs)
            # Negtive score function x reward
            loss = -m.log_prob(action) * reward # 核心
            # print(loss)
            loss.backward()

        self.optimizer.step()

    def save_model(self, path):
        torch.save(self.policy_net.state_dict(), path)

    def load_model(self, path):
        self.policy_net.load_state_dict(torch.load(path))

可以看到核心實現是以下幾句:

state = Variable(torch.from_numpy(state).float())
probs = self.policy_net(state)
m = Bernoulli(probs)
# Negtive score function x reward
loss = -m.log_prob(action) * reward # 核心
# print(loss)
loss.backward()

這裏採用的是伯努利分佈,二項分佈,舉個例子:

Example::
        >>> m = Bernoulli(torch.tensor([0.3]))
        >>> m.sample()  # 30% chance 1; 70% chance 0
        tensor([ 0.])

採樣結果是0或者1,1對應的概率是p,0對應概率是1-p。

爲神馬要用這個伯努利分佈呢?因爲這個這個問題是CartPole-v0 ,其動作空間只有0或1,所以這裏採用了Bernoulli,其他情況要使用不同的分佈才能滿足要求。

得到了採樣結果以後,就是用了第二節提到的REINFORCE的方法計算loss,進行loss反向傳播。

4. 總結

簡單介紹了以下如何使用,但並沒有深究背後的原理,這個系列會繼續更新,同時我也會繼續加強我的數學功底。

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章