Soccer-Environment-Multi-Agent-RL/model.py at master · AkshayS21/Soccer-Environment-Multi-Agent-RL · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71

# coding: utf-8

# In[3]:


import torch
import torch.nn as nn
import torch.nn.functional as f
import numpy as np


def hidden_init(layer):
    w_in = layer.weight.data.size()[0]
    lim = 1./np.sqrt(w_in)
    return (-lim,lim)


class Network(nn.Module):
    def __init__(self, input_dim ,h1, h2, h3, output_dim , actor = False):
        super(Network,self).__init__()

        self.input = input_dim
        self.h1 = h1
        self.h2 = h2
        self.h3 = h3
        self.output = output_dim
        self.actor = actor

        self.fc1 = nn.Linear(self.input,h1)
        self.fc2 = nn.Linear(self.h1,self.h2)
        self.fc3 = nn.Linear(self.h2,self.h3)
        self.fc4 = nn.Linear(self.h3, self.output)
        self.bn = nn.BatchNorm1d(self.h1)
        #self.reset_parameters()


    def reset_parameters(self):

        self.fc1.weight.data.uniform_(*hidden_init(self.fc1))
        self.fc2.weight.data.uniform_(*hidden_init(self.fc2))
        self.fc3.weight.data.uniform_(*hidden_init(self.fc3))

        self.fc4.weight.data.uniform_(-1e-3,1e3)


    def forward(self,x):

        if self.actor:
            if x.dim() == 1:
                x = torch.unsqueeze(x,0)

            x = f.relu(self.fc1(x))
            x = self.bn(x)
            x = f.relu(self.fc2(x))
            x = f.relu(self.fc3(x))
            return torch.tanh(self.fc4(x))
        else:

            x = f.relu(self.fc1(x))
            x = self.bn(x)
            x = f.relu(self.fc2(x))
            x = f.relu(self.fc3(x))
            return (self.fc4(x))