<gh_stars>1-10
# Code from https://github.com/Kaixhin/Rainbow
# -*- coding: utf-8 -*-
from __future__ import division
import math
import torch
from torch import nn
from torch.nn import functional as F


# Factorised NoisyLinear layer with bias
class NoisyLinear(nn.Module):
    def __init__(self, in_features, out_features, std_init=0.5):
        super(NoisyLinear, self).__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.std_init = std_init
        self.weight_mu = nn.Parameter(torch.empty(out_features, in_features))
        self.weight_sigma = nn.Parameter(torch.empty(out_features, in_features))
        self.register_buffer('weight_epsilon', torch.empty(out_features, in_features))
        self.bias_mu = nn.Parameter(torch.empty(out_features))
        self.bias_sigma = nn.Parameter(torch.empty(out_features))
        self.register_buffer('bias_epsilon', torch.empty(out_features))
        self.reset_parameters()
        self.reset_noise()

    def reset_parameters(self):
        mu_range = 1 / math.sqrt(self.in_features)
        self.weight_mu.data.uniform_(-mu_range, mu_range)
        self.weight_sigma.data.fill_(self.std_init / math.sqrt(self.in_features))
        self.bias_mu.data.uniform_(-mu_range, mu_range)
        self.bias_sigma.data.fill_(self.std_init / math.sqrt(self.out_features))

    def _scale_noise(self, size):
        x = torch.randn(size)
        return x.sign().mul_(x.abs().sqrt_())

    def reset_noise(self):
        epsilon_in = self._scale_noise(self.in_features)
        epsilon_out = self._scale_noise(self.out_features)
        self.weight_epsilon.copy_(epsilon_out.ger(epsilon_in))
        self.bias_epsilon.copy_(epsilon_out)

    def forward(self, input):
        if self.training:
          return F.linear(input, self.weight_mu + self.weight_sigma * self.weight_epsilon, self.bias_mu + self.bias_sigma * self.bias_epsilon)
        else:
          return F.linear(input, self.weight_mu, self.bias_mu)


# TODO Check hidsizes and pass them as params
class DQN(nn.Module):
    def __init__(self, args, state_size, action_space, hidsize1=128):
        super(DQN, self).__init__()
        self.atoms = args.atoms
        self.action_space = action_space
        # Not sure about initial layers, the original architecture uses conv layers
        # Need to choose the shared encoder
        self.fc_v = nn.Linear(state_size, hidsize1)
        self.fc_a = nn.Linear(state_size, hidsize1)
        self.fc_noisy_h_v = NoisyLinear(hidsize1, args.hidden_size, std_init=args.noisy_std)
        self.fc_noisy_h_a = NoisyLinear(hidsize1, args.hidden_size, std_init=args.noisy_std)
        self.fc_noisy_z_v = NoisyLinear(args.hidden_size, self.atoms, std_init=args.noisy_std)
        self.fc_noisy_z_a = NoisyLinear(args.hidden_size, action_space * self.atoms, std_init=args.noisy_std)

    def forward(self, x, log=False):

        # Value stream
        v = self.fc_noisy_z_v(F.relu(self.fc_noisy_h_v(F.relu(self.fc_v(x)))))
        # Advantage stream
        a = self.fc_noisy_z_a(F.relu(self.fc_noisy_h_a(F.relu(self.fc_a(x)))))
        v, a = v.view(-1, 1, self.atoms), a.view(-1, self.action_space, self.atoms)
        q = v + a - a.mean(1, keepdim=True)  # Combine streams
        if log:  # Use log softmax for numerical stability
            q = F.log_softmax(q, dim=2)  # Log probabilities with action over second dimension
        else:
            q = F.softmax(q, dim=2)  # Probabilities with action over second dimension
            
        return q # size(1, 2, 51) 51 as the atoms

    def reset_noise(self):
        for name, module in self.named_children():
            if 'noisy' in name:
                module.reset_noise()

# Code from https://github.com/Kaixhin/Rainbow
# -*- coding: utf-8 -*-
from __future__ import division
import math
import torch
from torch import nn
from torch.nn import functional as F


# Factorised NoisyLinear layer with bias
class NoisyLinear(nn.Module):
    def __init__(self, in_features, out_features, std_init=0.5):
        super(NoisyLinear, self).__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.std_init = std_init
        self.weight_mu = nn.Parameter(torch.empty(out_features, in_features))
        self.weight_sigma = nn.Parameter(torch.empty(out_features, in_features))
        self.register_buffer('weight_epsilon', torch.empty(out_features, in_features))
        self.bias_mu = nn.Parameter(torch.empty(out_features))
        self.bias_sigma = nn.Parameter(torch.empty(out_features))
        self.register_buffer('bias_epsilon', torch.empty(out_features))
        self.reset_parameters()
        self.reset_noise()

    def reset_parameters(self):
        mu_range = 1 / math.sqrt(self.in_features)
        self.weight_mu.data.uniform_(-mu_range, mu_range)
        self.weight_sigma.data.fill_(self.std_init / math.sqrt(self.in_features))
        self.bias_mu.data.uniform_(-mu_range, mu_range)
        self.bias_sigma.data.fill_(self.std_init / math.sqrt(self.out_features))

    def _scale_noise(self, size):
        x = torch.randn(size)
        return x.sign().mul_(x.abs().sqrt_())

    def reset_noise(self):
        epsilon_in = self._scale_noise(self.in_features)
        epsilon_out = self._scale_noise(self.out_features)
        self.weight_epsilon.copy_(epsilon_out.ger(epsilon_in))
        self.bias_epsilon.copy_(epsilon_out)

    def forward(self, input):
        if self.training:
          return F.linear(input, self.weight_mu + self.weight_sigma * self.weight_epsilon, self.bias_mu + self.bias_sigma * self.bias_epsilon)
        else:
          return F.linear(input, self.weight_mu, self.bias_mu)


# TODO Check hidsizes and pass them as params
class DQN(nn.Module):
    def __init__(self, args, state_size, action_space, hidsize1=128):
        super(DQN, self).__init__()
        self.atoms = args.atoms
        self.action_space = action_space
        # Not sure about initial layers, the original architecture uses conv layers
        # Need to choose the shared encoder
        self.fc_v = nn.Linear(state_size, hidsize1)
        self.fc_a = nn.Linear(state_size, hidsize1)
        self.fc_noisy_h_v = NoisyLinear(hidsize1, args.hidden_size, std_init=args.noisy_std)
        self.fc_noisy_h_a = NoisyLinear(hidsize1, args.hidden_size, std_init=args.noisy_std)
        self.fc_noisy_z_v = NoisyLinear(args.hidden_size, self.atoms, std_init=args.noisy_std)
        self.fc_noisy_z_a = NoisyLinear(args.hidden_size, action_space * self.atoms, std_init=args.noisy_std)

    def forward(self, x, log=False):

        # Value stream
        v = self.fc_noisy_z_v(F.relu(self.fc_noisy_h_v(F.relu(self.fc_v(x)))))
        # Advantage stream
        a = self.fc_noisy_z_a(F.relu(self.fc_noisy_h_a(F.relu(self.fc_a(x)))))
        v, a = v.view(-1, 1, self.atoms), a.view(-1, self.action_space, self.atoms)
        q = v + a - a.mean(1, keepdim=True)  # Combine streams
        if log:  # Use log softmax for numerical stability
            q = F.log_softmax(q, dim=2)  # Log probabilities with action over second dimension
        else:
            q = F.softmax(q, dim=2)  # Probabilities with action over second dimension
            
        return q # size(1, 2, 51) 51 as the atoms

    def reset_noise(self):
        for name, module in self.named_children():
            if 'noisy' in name:
                module.reset_noise()
# Code from https://github.com/Kaixhin/Rainbow # -*- coding: utf-8 -*- # Factorised NoisyLinear layer with bias # TODO Check hidsizes and pass them as params # Not sure about initial layers, the original architecture uses conv layers # Need to choose the shared encoder # Value stream # Advantage stream # Combine streams # Use log softmax for numerical stability # Log probabilities with action over second dimension # Probabilities with action over second dimension # size(1, 2, 51) 51 as the atoms