ContinuousControl-D4PG/model.py at master · pablobd/ContinuousControl-D4PG · 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
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
import numpy as np

import torch
import torch.nn as nn
import torch.nn.functional as F

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

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

class Actor(nn.Module):
    " Actor (Policy) Model - Neural net to decide what action the agent must take "

    def __init__(self, action_size, state_size, seed, hidden_layers, init_weights):
        """
        Initialize parameters and build model.

        Params
        ======
            state_size (int): Dimension of each state
            action_size (int): Dimension of each action
            hidden_layers (int list): Number of hidden layers and nodes in each hidden layer
            seed (int): Random seed
        """
        super().__init__()
        self.seed = torch.manual_seed(seed)

        # report dimensions of hidden layer
        print("Hidden layers Actor: ", hidden_layers)

        # initial layer
        self.batch_norm = nn.ModuleList([nn.BatchNorm1d(state_size)])
        self.hidden_layers = nn.ModuleList([nn.Linear(state_size, hidden_layers[0])])

        # hidden layers
        layer_sizes = zip(hidden_layers[:-1], hidden_layers[1:])

        self.hidden_layers.extend([nn.Linear(h1, h2) for h1, h2 in layer_sizes])
        self.batch_norm.extend([nn.BatchNorm1d(h) for h in hidden_layers[:-1]])

        # final layer
        self.output = nn.Linear(hidden_layers[-1], action_size)


        # send networks to device
        for linear in self.hidden_layers:
            linear.to(device)

        self.output.to(device)

        if init_weights:
            print("initializing weights...")
            self.initialize_weights()

    def initialize_weights(self):
        " Initialize weights of layers "

        for layer in self.hidden_layers:
            layer.weight.data.uniform_(*hidden_init(layer))

        self.output.weight.data.uniform_(-3e-3, 3e-3)


    def forward(self, x):
        """
        Forward network that maps state -> action

        Params
        ======
            x (tensor): State vector
        """

        # forward through each layer in `hidden_layers`, with ReLU activation
        layers = zip(self.batch_norm, self.hidden_layers)
        for batch, linear in layers:
            x = batch(x)
            x = F.leaky_relu(linear(x))

        x = self.output(x)

        # forward final layer with tanh activation (-1, 1)
        return torch.tanh(x)


class Critic(nn.Module):
    " Critic (Value) Model - Neural net to estimate the total expected episodic return associated to one action in a given state "

    def __init__(self, action_size, state_size, seed, hidden_layers, init_weights):
        """
        Initialize parameters and build model.

        Params
        ======
            state_size (int): Dimension of each state
            action_size (int): Dimension of each action
            hidden_layers (int list): Number of hidden layers and nodes in each hidden layer
            seed (int): Random seed
        """
        super().__init__()
        self.seed = torch.manual_seed(seed)

        print("Hidden layers Critic: ", hidden_layers)

        # initial layer
        self.batch = nn.BatchNorm1d(state_size)
        self.hidden_layers = nn.ModuleList([nn.Linear(state_size, hidden_layers[0])])

        # hidden layers
        hidden_layers[0] += action_size
        layer_sizes = zip(hidden_layers[:-1], hidden_layers[1:])

        self.hidden_layers.extend([nn.Linear(h1, h2) for h1, h2 in layer_sizes])

        # final layer
        self.output = nn.Linear(hidden_layers[-1], 1)

        # send networks to device
        for linear in self.hidden_layers:
            linear.to(device)

        self.output.to(device)

        if init_weights:
            print("initializing weights...")
            self.initialize_weights()

    def initialize_weights(self):
        " Initialize weights of layers "

        for layer in self.hidden_layers:
            layer.weight.data.uniform_(*hidden_init(layer))

        self.output.weight.data.uniform_(-3e-3, 3e-3)

    def forward(self, states, actions):
        """
        Forward network that maps state -> action

        Params
        ======
            states (tensor): State vector
            actions (tensor): Action vector
        """

        # forward through first layer
        x = self.batch(states)
        x = F.leaky_relu(self.hidden_layers[0](x))

        # concatenate output of first layer and action vector
        x = torch.cat((x, actions), dim = 1)

        # forward through each layer in `hidden_layers`, with Leaky ReLU activation
        for linear in self.hidden_layers[1:]:
            x = F.leaky_relu(linear(x))

        return self.output(x)