11.06课堂作业2：PPO算法详细步骤

def PPO_algorithm():
    # 初始化
    policy_network = PolicyNetwork()
    value_network = ValueNetwork()
    optimizer = Adam(list(policy_network.parameters()) + list(value_network.parameters()))
    
    for iteration in range(max_iterations):
        # 数据收集阶段
        states, actions, rewards, log_probs, values, dones = [], [], [], [], [], []
        state = env.reset()
        
        for t in range(batch_size):
            # 策略执行
            action_dist = policy_network(state)
            action = action_dist.sample()
            log_prob = action_dist.log_prob(action)
            value = value_network(state)
            
            # 环境交互
            next_state, reward, done, _ = env.step(action)
            
            # 存储数据
            states.append(state)
            actions.append(action)
            rewards.append(reward)
            log_probs.append(log_prob)
            values.append(value)
            dones.append(done)
            
            state = next_state
            if done:
                state = env.reset()
        
        # 计算优势函数
        advantages, returns = compute_gae(rewards, values, dones, gamma, lambda_)
        
        # 策略更新阶段
        old_log_probs = log_probs.clone()
        
        for epoch in range(update_epochs):
            # 随机打乱数据
            indices = torch.randperm(batch_size)
            
            for start in range(0, batch_size, mini_batch_size):
                end = start + mini_batch_size
                batch_indices = indices[start:end]
                
                # 计算新的策略概率和价值
                batch_states = states[batch_indices]
                batch_actions = actions[batch_indices]
                batch_old_log_probs = old_log_probs[batch_indices]
                batch_advantages = advantages[batch_indices]
                batch_returns = returns[batch_indices]
                
                action_dist = policy_network(batch_states)
                new_log_probs = action_dist.log_prob(batch_actions)
                new_values = value_network(batch_states)
                
                # 计算比率
                ratio = torch.exp(new_log_probs - batch_old_log_probs)
                
                # 计算策略损失
                surr1 = ratio * batch_advantages
                surr2 = torch.clamp(ratio, 1 - epsilon, 1 + epsilon) * batch_advantages
                policy_loss = -torch.min(surr1, surr2).mean()
                
                # 计算价值损失
                value_loss = F.mse_loss(new_values, batch_returns)
                
                # 计算熵损失
                entropy_loss = -action_dist.entropy().mean()
                
                # 总损失
                total_loss = policy_loss + c1 * value_loss + c2 * entropy_loss
                
                # 更新参数
                optimizer.zero_grad()
                total_loss.backward()
                optimizer.step()
        
        # 记录训练统计信息
        log_training_stats(iteration, policy_loss, value_loss, entropy_loss)
        
        # 评估模型
        if iteration % eval_freq == 0:
            avg_reward = evaluate_model(policy_network, env)
            print(f"Iteration {iteration}: Average Reward = {avg_reward:.2f}")


# GAE计算函数
def compute_gae(rewards, values, dones, gamma, lambda_):
    advantages = []
    returns = []
    gae = 0
    
    for t in reversed(range(len(rewards))):
        if t == len(rewards) - 1:
            next_value = 0
        else:
            next_value = values[t + 1]
        
        delta = rewards[t] + gamma * next_value * (1 - dones[t]) - values[t]
        gae = delta + gamma * lambda_ * (1 - dones[t]) * gae
        advantages.insert(0, gae)
        returns.insert(0, gae + values[t])
    
    return torch.tensor(advantages), torch.tensor(returns)