آموزش یادگیری تقویتی چندعامله (MARL) با QMIX برای بهینهسازی مسیریابی رباتها

import torch
import torch.nn as nn
import torch.nn.functional as F

class AgentDRQN(nn.Module):
    """
    Decentralized Recurrent Q-Network per agent.
    Uses GRU to encode local action-observation history tau^a.
    """
    def __init__(self, obs_dim: int, n_actions: int, hidden_dim: int = 64):
        super().__init__()
        self.fc1 = nn.Linear(obs_dim, hidden_dim)
        self.rnn = nn.GRUCell(hidden_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, n_actions)

    def forward(self, obs: torch.Tensor, hidden_state: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        x = F.relu(self.fc1(obs))
        h_next = self.rnn(x, hidden_state)
        q_values = self.fc2(h_next)
        return q_values, h_next

class QMixMonotonicMixer(nn.Module):
    """
    Centralized Monotonic Mixing Network powered by Hypernetworks.
    Enforces dQ_tot / dQ_a >= 0 via absolute value weight constraints.
    """
    def __init__(self, n_agents: int, state_dim: int, mixing_embed_dim: int = 32, hypernet_embed: int = 64):
        super().__init__()
        self.n_agents = n_agents
        self.state_dim = state_dim
        self.embed_dim = mixing_embed_dim

        # Hypernetworks generating strictly positive mixing weights W1 and W2
        self.hyper_w1 = nn.Sequential(
            nn.Linear(state_dim, hypernet_embed),
            nn.ReLU(),
            nn.Linear(hypernet_embed, n_agents * mixing_embed_dim)
        )
        self.hyper_b1 = nn.Linear(state_dim, mixing_embed_dim)

        self.hyper_w2 = nn.Sequential(
            nn.Linear(state_dim, hypernet_embed),
            nn.ReLU(),
            nn.Linear(hypernet_embed, mixing_embed_dim)
        )
        # State-dependent final bias V(s) (unconstrained sign)
        self.state_value_v = nn.Sequential(
            nn.Linear(state_dim, mixing_embed_dim),
            nn.ReLU(),
            nn.Linear(mixing_embed_dim, 1)
        )

    def forward(self, agent_qs: torch.Tensor, states: torch.Tensor) -> torch.Tensor:
        """
        agent_qs: (batch_size, seq_len, n_agents)
        states:   (batch_size, seq_len, state_dim)
        """
        bs, seq_len, _ = agent_qs.size()
        agent_qs = agent_qs.view(-1, 1, self.n_agents)
        states = states.reshape(-1, self.state_dim)

        # Enforce monotonicity constraint: W1 >= 0
        w1 = torch.abs(self.hyper_w1(states)).view(-1, self.n_agents, self.embed_dim)
        b1 = self.hyper_b1(states).view(-1, 1, self.embed_dim)
        hidden = F.elu(torch.bmm(agent_qs, w1) + b1)

        # Enforce monotonicity constraint: W2 >= 0
        w2 = torch.abs(self.hyper_w2(states)).view(-1, self.embed_dim, 1)
        v = self.state_value_v(states).view(-1, 1, 1)

        q_tot = torch.bmm(hidden, w2) + v
        return q_tot.view(bs, seq_len, 1)

def compute_qmix_td_loss(
    chosen_action_qs: torch.Tensor,
    target_max_qs: torch.Tensor,
    states: torch.Tensor,
    next_states: torch.Tensor,
    rewards: torch.Tensor,
    terminated: torch.Tensor,
     masks: torch.Tensor,
    mixer: QMixMonotonicMixer,
    target_mixer: QMixMonotonicMixer,
    gamma: float = 0.99
) -> torch.Tensor:
    q_tot = mixer(chosen_action_qs, states)
    with torch.no_grad():
        target_q_tot = target_mixer(target_max_qs, next_states)
        y_dqn = rewards + gamma * (1.0 - terminated) * target_q_tot

    td_error = (q_tot - y_dqn) * masks
    return (td_error ** 2).sum() / masks.sum()