q_mixer.py

单调 mixer

原文件 · 50 行 · 2031 字节

algorithms/qmix/algorithm/q_mixer.py
SHA-256: 4407a88f17e2c6af0ac2c7e5eaa0a4e40bdc5b15be5e243194155f0390419889

1"""Monotonic mixing network used by the local QMIX implementation."""23import torch4import torch.nn as nn5import torch.nn.functional as F678class QMixer(nn.Module):9    """Map per-agent utilities and the centralized state to a joint Q value.1011    Absolute hypernetwork weights enforce the monotonicity constraint from12    QMIX, so decentralized greedy actions maximize the centralized value.13    Inputs use ``[..., num_agents]`` and ``[..., state_dim]`` shapes.14    """1516    def __init__(self, state_dim, num_agents, hidden_dim=32, hypernet_hidden_dim=64):17        super().__init__()18        self.state_dim = int(state_dim)19        self.num_agents = int(num_agents)20        self.hidden_dim = int(hidden_dim)21        hypernet_hidden_dim = int(hypernet_hidden_dim)22        self.hyper_w1 = nn.Sequential(23            nn.Linear(self.state_dim, hypernet_hidden_dim),24            nn.ReLU(),25            nn.Linear(hypernet_hidden_dim, self.num_agents * self.hidden_dim),26        )27        self.hyper_b1 = nn.Linear(self.state_dim, self.hidden_dim)28        self.hyper_w2 = nn.Sequential(29            nn.Linear(self.state_dim, hypernet_hidden_dim),30            nn.ReLU(),31            nn.Linear(hypernet_hidden_dim, self.hidden_dim),32        )33        self.hyper_b2 = nn.Sequential(34            nn.Linear(self.state_dim, hypernet_hidden_dim),35            nn.ReLU(),36            nn.Linear(hypernet_hidden_dim, 1),37        )3839    def forward(self, agent_qs, states):40        prefix = agent_qs.shape[:-1]41        flat_qs = agent_qs.reshape(-1, 1, self.num_agents)42        flat_states = states.reshape(-1, self.state_dim)43        w1 = self.hyper_w1(flat_states).abs().view(44            -1, self.num_agents, self.hidden_dim45        )46        b1 = self.hyper_b1(flat_states).view(-1, 1, self.hidden_dim)47        hidden = F.elu(torch.bmm(flat_qs, w1) + b1)48        w2 = self.hyper_w2(flat_states).abs().view(-1, self.hidden_dim, 1)49        b2 = self.hyper_b2(flat_states).view(-1, 1, 1)50        return (torch.bmm(hidden, w2) + b2).view(*prefix, 1)