Search code examples
modelreinforcement-learning

Output the weights from a Pytorch model


Here a very basic model :

class LinearDeepQNetwork(nn.Module):
    def __init__(self, lr, n_actions, input_dims):
        super(LinearDeepQNetwork, self).__init__()

        self.fc1 = nn.Linear(*input_dims, 128)
        self.fc2 = nn.Linear(128, n_actions)

        self.optimizer = optim.Adam(self.parameters(), lr=lr)
        self.loss = nn.MSELoss()
        self.device = T.device('cuda:0' if T.cuda.is_available() else 'cpu')
        self.to(self.device)

    def forward(self, state):
        layer1 = F.relu(self.fc1(state))
        actions = self.fc2(layer1)

        return actions

Be aware that I am using Pytorch, not Keras or Tensorflow. In my Agent() class, I instantiate self.Q_eval = LinearDeepQNetwork(self.lr, self.n_actions, self.input_dims). Once I have trained my agent for several episodes, I need to output the weights of self.Q_eval. How can I do that?


Solution

  • I needed to inject the weights from Q_eval network to Q_next network. I made the following function :

    def replace_target_network(self):
            self.Q_next.load_state_dict(self.Q_eval.state_dict())
            self.Q_next.eval()
    

    Answer I can get the weights with Q_eval.state_dict().