A2C python实现

A2C (Advantage Actor-Critic) 是一种强化学习算法，它结合了 Actor-Critic 和 Advantage 的优点，可以有效地解决强化学习中的稳定性问题。以下是一个简单的 A2C 的 Python 实现：

import gym
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.distributions import Categorical

# 定义 Actor 网络
class Actor(nn.Module):
    def __init__(self, state_dim, action_dim, hidden_dim):
        super(Actor, self).__init__()
        self.fc1 = nn.Linear(state_dim, hidden_dim)
        self.action_head = nn.Linear(hidden_dim, action_dim)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        action_probs = F.softmax(self.action_head(x), dim=-1)
        return action_probs

# 定义 Critic 网络
class Critic(nn.Module):
    def __init__(self, state_dim, hidden_dim):
        super(Critic, self).__init__()
        self.fc1 = nn.Linear(state_dim, hidden_dim)
        self.value_head = nn.Linear(hidden_dim, 1)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        state_value = self.value_head(x)
        return state_value

# 定义 A2C 算法
class A2C:
    def __init__(self, state_dim, action_dim, hidden_dim, lr, gamma):
        self.actor = Actor(state_dim, action_dim, hidden_dim)
        self.critic = Critic(state_dim, hidden_dim)
        self.optimizer_actor = optim.Adam(self.actor.parameters(), lr=lr)
        self.optimizer_critic = optim.Adam(self.critic.parameters(), lr=lr)
        self.gamma = gamma

    def train(self, env, episodes):
        for i in range(episodes):
            state = env.reset()
            rewards = []
            log_probs = []
            values = []
            while True:
                # 选择动作
                state = torch.from_numpy(state).float().unsqueeze(0)
                action_probs = self.actor(state)
                dist = Categorical(action_probs)
                action = dist.sample()
                log_prob = dist.log_prob(action)

                # 执行动作
                next_state, reward, done, _ = env.step(action.item())
                rewards.append(reward)
                log_probs.append(log_prob)
                values.append(self.critic(state))

                state = next_state

                if done:
                    # 计算返回值
                    returns = [0]
                    for r in reversed(rewards):
                        returns[0] = r + self.gamma * returns[0]
                        returns.insert(0, returns[0])
                    returns = returns[:-1]

                    # 计算 Advantage
                    advantages = [ret - value.item() for ret, value in zip(returns, values)]

                    # 计算 actor 和 critic 的损失函数
                    actor_loss = 0
                    critic_loss = 0
                    for log_prob, advantage, value in zip(log_probs, advantages, values):
                        actor_loss += -log_prob * advantage
                        critic_loss += F.smooth_l1_loss(value, torch.tensor([returns]))

                    # 更新网络参数
                    self.optimizer_actor.zero_grad()
                    actor_loss.backward()
                    self.optimizer_actor.step()

                    self.optimizer_critic.zero_grad()
                    critic_loss.backward()
                    self.optimizer_critic.step()

                    break

# 测试 A2C 算法
env = gym.make('CartPole-v0')
a2c = A2C(state_dim=4, action_dim=2, hidden_dim=128, lr=0.001, gamma=0.99)
a2c.train(env, episodes=1000)

这是一个简单的 A2C 实现，可以在 Gym 的 CartPole-v0 环境上进行训练和测试。在训练过程中，会输出每个 episode 的平均奖励和最后一个 episode 的奖励。