理论部分以及完整代码参看之前的博客：https://blog.csdn.net/qq_47997583/article/details/124506650
本文章介绍的是策略梯度算法中的REINFORCE实现
上图为算法流程图，总体来说代码实现中，先实现一个episode然后从后往前计算回报，损失函数是负的回报乘于log的该状态下采取该动作的概率。每个状态动作对对应算一次loss，然后反向传播计算梯度。最后整个episode完之后进行梯度下降。
在代码实现中我们需要是实现两个类PolicyNet，REINFORCE以及主函数部分。
主函数部分
for i in range(10):
with tqdm(total=int(num_episodes / 10), desc='Iteration %d' % i) as pbar:
for i_episode in range(int(num_episodes / 10)):
episode_return = 0
transition_dict = {
'states': [],
'actions': [],
'next_states': [],
'rewards': [],
'dones': []
}
state = env.reset()
done = False
while not done:
action = agent.take_action(state)
next_state, reward, done, _ = env.step(action)
transition_dict['states'].append(state)
transition_dict['actions'].append(action)
transition_dict['next_states'].append(next_state)
transition_dict['rewards'].append(reward)
transition_dict['dones'].append(done)
state = next_state
episode_return += reward
return_list.append(episode_return)
agent.update(transition_dict)
if (i_episode + 1) % 10 == 0:
pbar.set_postfix({
'episode':
'%d' % (num_episodes / 10 * i + i_episode + 1),
'return':
'%.3f' % np.mean(return_list[-10:])
})
pbar.update(1)
首先从主函数部分看，一共是训练1000个episode，分为10个iteration，每个100个episode。每个episode我们需要先初始化一个回合总奖励和一个字典记录整个回合每个时间步的五元组信息；之后初始化状态state和done，然后进行采样直到回合结束，采样的动作通过REINFORCE类实例化的agent的take_action方法获得，然后将动作传入step函数获得四元组，将五元组传入字典中；回合结束后将累计回合奖励传入结果列表，之后通过agent的update方法更新参数。
PolicyNet
class PolicyNet(torch.nn.Module):
def __init__(self, state_dim, hidden_dim, action_dim):
super(PolicyNet, self).__init__()
self.fc1 = torch.nn.Linear(state_dim, hidden_dim)
self.fc2 = torch.nn.Linear(hidden_dim, action_dim)
def forward(self, x):
x = F.relu(self.fc1(x))
return F.softmax(self.fc2(x), dim=1)
# 0是对列做归一化，1是对行做归一化
REINFORCE
class REINFORCE:
def __init__(self, state_dim, hidden_dim, action_dim, learning_rate, gamma,
device):
self.policy_net = PolicyNet(state_dim, hidden_dim,
action_dim).to(device)
self.optimizer = torch.optim.Adam(self.policy_net.parameters(),
lr=learning_rate) # 使用Adam优化器
self.gamma = gamma # 折扣因子
self.device = device
def take_action(self, state):
# 根据动作概率分布随机采样
state = torch.tensor([state], dtype=torch.float).to(self.device)
# 1*4
probs = self.policy_net(state)
# 1*2
action_dist = torch.distributions.Categorical(probs)
action = action_dist.sample()
return action.item()
def update(self, transition_dict):
reward_list = transition_dict['rewards']
state_list = transition_dict['states']
action_list = transition_dict['actions']
G = 0
self.optimizer.zero_grad()
for i in reversed(range(len(reward_list))):
# 从最后一步算起
reward = reward_list[i]
state = torch.tensor([state_list[i]],
# 1*4
dtype=torch.float).to(self.device)
action = torch.tensor([action_list[i]]).view(-1, 1).to(self.device)
# 1*1
log_prob = torch.log(self.policy_net(state).gather(1, action))
# 1*1
G = self.gamma * G + reward
loss = -log_prob * G # 每一步的损失函数
loss.backward() # 反向传播计算梯度
self.optimizer.step() # 梯度下降
在实现REINFORCE类时我们需要实现两个方法，一个是take_action,一个是update。take_action作用是传入state到神经网络获得两个动作的概率分布，然后依据概率分布进行动作的抽样，update则是本算法的核心部分，我们需要传入包含回合的五元组数据的字典transition_dict，拿出来其中的s,a,r列表。之后首先初始化累计奖励G为0并且将梯度清零，然后执行循环，循环的次数为列表的长度，每次循环从列表末尾往前遍历：获得reward，state，action，log_prob通过将state传入神经网络获得两个动作的概率，然后根据action索引得到在本次时间步的π(a|s)，然后计算其log值。由于REINFORCE算法是计算当前时间之后的累计奖励作为回报，因此 G 的更新方式为 G = self.gamma * G + reward，将G乘负的对数log_prob即为损失函数，然后进行反向传播进行梯度累计。最后循环结束也就是episode结束，再进行梯度下降更新神经网络参数。
