【導語】:在深度強化學習第四篇中,講了Policy Gradient的理論。通過最終推導得到的公式,本文用PyTorch簡單實現以下,並且盡可能搞清楚torch.distribution的使用方法。代碼參考了LeeDeepRl-Notes中的實現。
1. 復習
\(\theta\)代表模型的參數,第一行公式代表了模型進行更新的方法,\(\eta\) 代表的是學習率。
第二行是推導得到的,和CrossEntropy可以對照着理解記憶。
2. Torch.Distributions
distributions包主要是實現了參數化的概率分布和采樣函數。參數化是為了讓模型能夠具有反向傳播的能力,這樣才可以用隨機梯度下降的方法來進行優化。隨機采樣的話沒辦法直接反向傳播,有兩個方法,REINFORCE和pathwise derivative estimator。
Torch中提供兩個方法,sample()和log_prob(),就可以實現REINFORCE
\(\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. 總結
簡單介紹了以下如何使用,但並沒有深究背后的原理,這個系列會繼續更新,同時我也會繼續加強我的數學功底。