多智能体编队actor-critic pytorch代码

以下是一个简单的多智能体编队actor-critic PyTorch代码示例：

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np

# 定义Actor神经网络
class Actor(nn.Module):
    def __init__(self, state_size, action_size):
        super(Actor, self).__init__()
        self.fc1 = nn.Linear(state_size, 128)
        self.fc2 = nn.Linear(128, 64)
        self.fc3 = nn.Linear(64, action_size)
        self.relu = nn.ReLU()
        self.softmax = nn.Softmax(dim=-1)

    def forward(self, state):
        x = self.relu(self.fc1(state))
        x = self.relu(self.fc2(x))
        x = self.softmax(self.fc3(x))
        return x

# 定义Critic神经网络
class Critic(nn.Module):
    def __init__(self, state_size):
        super(Critic, self).__init__()
        self.fc1 = nn.Linear(state_size, 128)
        self.fc2 = nn.Linear(128, 64)
        self.fc3 = nn.Linear(64, 1)
        self.relu = nn.ReLU()

    def forward(self, state):
        x = self.relu(self.fc1(state))
        x = self.relu(self.fc2(x))
        x = self.fc3(x)
        return x

# 定义Actor-Critic算法
class ActorCritic:
    def __init__(self, state_size, action_size, lr_actor=1e-4, lr_critic=1e-3, gamma=0.99):
        self.actor = Actor(state_size, action_size)
        self.critic = Critic(state_size)
        self.optimizer_actor = optim.Adam(self.actor.parameters(), lr=lr_actor)
        self.optimizer_critic = optim.Adam(self.critic.parameters(), lr=lr_critic)
        self.gamma = gamma

    def select_action(self, state):
        state = torch.FloatTensor(state)
        action_probs = self.actor.forward(state)
        action = torch.multinomial(action_probs, 1)
        return action.item()

    def update(self, rewards, states, next_states, actions, done):
        # 计算critic的loss
        rewards = torch.FloatTensor(rewards)
        states = torch.FloatTensor(states)
        next_states = torch.FloatTensor(next_states)
        actions = torch.LongTensor(actions)

        td_target = rewards + (1 - done) * self.gamma * self.critic(next_states).squeeze()
        td_error = td_target - self.critic(states).squeeze()
        critic_loss = td_error.pow(2).mean()

        # 更新critic网络
        self.optimizer_critic.zero_grad()
        critic_loss.backward()
        self.optimizer_critic.step()

        # 计算actor的loss
        action_probs = self.actor(states)
        log_probs = torch.log(torch.gather(action_probs, 1, actions.view(-1, 1)))
        actor_loss = -(log_probs * td_error.detach()).mean()

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

# 定义环境
class Environment:
    def __init__(self, num_agents, state_size, action_size):
        self.num_agents = num_agents
        self.state_size = state_size
        self.action_size = action_size

    def reset(self):
        self.states = np.zeros((self.num_agents, self.state_size))
        self.rewards = np.zeros(self.num_agents)
        self.done = np.zeros(self.num_agents, dtype=bool)
        self.total_reward = np.zeros(self.num_agents)
        return self.states

    def step(self, actions):
        for i in range(self.num_agents):
            if not self.done[i]:
                self.states[i] += np.random.rand(self.state_size)  # 模拟环境
                self.rewards[i] = np.random.rand()  # 模拟奖励
                self.total_reward[i] += self.rewards[i]
                if np.random.rand() < 0.1:  # 模拟终止
                    self.done[i] = True
        next_states = self.states
        return next_states, self.rewards, self.done

# 训练
env = Environment(num_agents=2, state_size=4, action_size=2)
ac = ActorCritic(state_size=4, action_size=2)

num_episodes = 1000
for i in range(num_episodes):
    states = env.reset()
    done = False
    while not done.all():
        actions = [ac.select_action(state) for state in states]
        next_states, rewards, done = env.step(actions)
        ac.update(rewards, states, next_states, actions, done)
        states = next_states
    print("Episode {}, Total reward: {}".format(i, env.total_reward))

请注意，此代码示例仅用于说明多智能体编队actor-critic算法的实现方式，并不是一个完整、可用的算法。实际应用中，需要根据具体问题进行修改和优化。