agent_q_function.py

个体 Q 网络

原文件 · 30 行 · 1121 字节

algorithms/qmix/algorithm/agent_q_function.py
SHA-256: 4912bb8d794bb51461723a8bb1e20891cfcb0aa1525f3d7f52c1d58bb5e09e00

1"""Compatibility wrapper around the maintained recurrent individual-Q net."""23import torch45from algorithms.qmix.algorithm.QMixPolicy import RecurrentAgentQNetwork678class AgentQFunction(RecurrentAgentQNetwork):9    def __init__(self, args, input_dim, act_dim, device):10        hidden_dim = int(11            getattr(args, "qmix_agent_hidden_dim", getattr(args, "hidden_size", 128))12        )13        super().__init__(input_dim, act_dim, hidden_dim)14        self.device = device15        self.to(device)1617    def forward(self, obs, rnn_states):18        obs = torch.as_tensor(obs, dtype=torch.float32, device=self.device)19        hidden = torch.as_tensor(rnn_states, dtype=torch.float32, device=self.device)20        if obs.ndim == 2:21            if hidden.ndim == 3:22                hidden = hidden[0]23            return self.forward_step(obs, hidden)24        if hidden.ndim == 3:25            hidden = hidden[0]26        outputs = []27        for step in range(obs.shape[0]):28            q_values, hidden = self.forward_step(obs[step], hidden)29            outputs.append(q_values)30        return torch.stack(outputs), hidden