#!/usr/bin/env python
# coding=utf-8
"""
@author: <NAME>
@create time: 2019-12-07 19:53
@edit time: 2021-01-25 14:06
@desc: DQN 使用的网络结构，其特点是输出各个离散动作值的value
"""


import torch
import torch.nn as nn


class DQNNet(nn.Module):
    def __init__(self, n_states, n_actions):
        """
        定义隐藏层和输出层参数
        @param n_obs: number of observations
        @param n_actions: number of actions
        """
        super(DQNNet, self).__init__()
        self.seq = nn.Sequential(
            nn.Linear(n_states, 128),
            nn.ReLU(),
            nn.Linear(128, 256),
            nn.ReLU(),
            nn.Linear(256, 256),
            nn.ReLU(),
            nn.Linear(256, n_actions),
        )

    def forward(self, x):
        """
        定义网络结构: 第一层网络->ReLU激活->输出层->softmax->输出
        """
        action_values = self.seq(x)
        return action_values


# DDQN 用的 Net 和 DQN 是一致的
DoubleDQNNet = DDQNNet = DQNNet


class DuelingDQNNet(nn.Module):
    """Double DQN神经网络"""

    def __init__(self, n_states, n_actions):
        """
        定义隐藏层和输出层参数
        @param n_obs: number of observations
        @param n_actions: number of actions
        @param n_neurons: number of neurons for the hidden layer
        """
        super(DuelingDQNNet, self).__init__()
        self.net_state = nn.Sequential(
            nn.Linear(n_states, 128), nn.ReLU(),
            nn.Linear(128, 256), nn.ReLU(),
        )
        self.net_val = nn.Sequential(
            nn.Linear(256, 256), nn.ReLU(),
            nn.Linear(256, 1),   # Q value
            nn.ReLU()
        )
        self.net_adv = nn.Sequential(
            nn.Linear(256, 256), nn.ReLU(),
            nn.Linear(256, n_actions),
            nn.ReLU()
        )

    def forward(self, x):
        x = self.net_state(x)
        q_val = self.net_val(x)
        q_adv = self.net_adv(x)
        # dueling Q value
        return q_val + q_adv - q_adv.mean(dim=1, keepdim=True)


# D3QN的网络和Dueling保持一致
D3QNNet = DuelingDQNNet

#!/usr/bin/env python
# coding=utf-8
"""
@author: <NAME>
@create time: 2019-12-07 19:53
@edit time: 2021-01-25 14:06
@desc: DQN 使用的网络结构，其特点是输出各个离散动作值的value
"""


import torch
import torch.nn as nn


class DQNNet(nn.Module):
    def __init__(self, n_states, n_actions):
        """
        定义隐藏层和输出层参数
        @param n_obs: number of observations
        @param n_actions: number of actions
        """
        super(DQNNet, self).__init__()
        self.seq = nn.Sequential(
            nn.Linear(n_states, 128),
            nn.ReLU(),
            nn.Linear(128, 256),
            nn.ReLU(),
            nn.Linear(256, 256),
            nn.ReLU(),
            nn.Linear(256, n_actions),
        )

    def forward(self, x):
        """
        定义网络结构: 第一层网络->ReLU激活->输出层->softmax->输出
        """
        action_values = self.seq(x)
        return action_values


# DDQN 用的 Net 和 DQN 是一致的
DoubleDQNNet = DDQNNet = DQNNet


class DuelingDQNNet(nn.Module):
    """Double DQN神经网络"""

    def __init__(self, n_states, n_actions):
        """
        定义隐藏层和输出层参数
        @param n_obs: number of observations
        @param n_actions: number of actions
        @param n_neurons: number of neurons for the hidden layer
        """
        super(DuelingDQNNet, self).__init__()
        self.net_state = nn.Sequential(
            nn.Linear(n_states, 128), nn.ReLU(),
            nn.Linear(128, 256), nn.ReLU(),
        )
        self.net_val = nn.Sequential(
            nn.Linear(256, 256), nn.ReLU(),
            nn.Linear(256, 1),   # Q value
            nn.ReLU()
        )
        self.net_adv = nn.Sequential(
            nn.Linear(256, 256), nn.ReLU(),
            nn.Linear(256, n_actions),
            nn.ReLU()
        )

    def forward(self, x):
        x = self.net_state(x)
        q_val = self.net_val(x)
        q_adv = self.net_adv(x)
        # dueling Q value
        return q_val + q_adv - q_adv.mean(dim=1, keepdim=True)


# D3QN的网络和Dueling保持一致
D3QNNet = DuelingDQNNet

#!/usr/bin/env python # coding=utf-8 @author: <NAME> @create time: 2019-12-07 19:53 @edit time: 2021-01-25 14:06 @desc: DQN 使用的网络结构，其特点是输出各个离散动作值的value 定义隐藏层和输出层参数         @param n_obs: number of observations         @param n_actions: number of actions 定义网络结构: 第一层网络->ReLU激活->输出层->softmax->输出 # DDQN 用的 Net 和 DQN 是一致的 Double DQN神经网络 定义隐藏层和输出层参数         @param n_obs: number of observations         @param n_actions: number of actions         @param n_neurons: number of neurons for the hidden layer # Q value # dueling Q value # D3QN的网络和Dueling保持一致