Custom Model with custom submodels in it not updating

Khabbab_Zakaria: <p>Hi there,</p>
<p>I need to use a Reinforcement Learning Model where I need to use a RNN model. I wrote the code for the RNN like this:</p>
<pre><code class="lang-auto">class LSTMModel(nn.Module):
    def __init__(self, input_dim, hidden_dim, layer_dim, output_dim, dropout_prob, state_size):
        super(LSTMModel, self).__init__()

        # Defining the number of layers and the nodes in each layer
        self.hidden_dim = hidden_dim
        self.layer_dim = layer_dim
        ........
   def forward(self):
        .......
        return ....
rnn = LSTMModel(input_dim, hidden_dim, layer_dim, output_dim, dropout, state_size)
</code></pre>
<p>And I wrote the policy like this:</p>
<pre><code class="lang-auto">class Policy(nn.Module):
    """
    implements both actor and critic in one model
    """
    def __init__(self):
        super(Policy, self).__init__()
    ............
   def forward(self):
       ..............
policy = Policy()
</code></pre>
<p>Both of these are used in the model like this:</p>
<pre><code class="lang-auto">class Model(torch.nn.Module):
    def __init__(self, ann_private, rnn, policy):
        super(Model, self).__init__()
        self.rnn = rnn
        self.policy = policy

    def forward(self, private_input, state):
          .....
model = Model(rnn, policy)
</code></pre>
<p>The optimizer is written as optimizer = optim.Adam(model.parameters(), lr=1e-3)</p>
<p>The loss is calculated in a traditional actor-critic way. In each ‘subepisode’, we calculate the states, actions, rewards, next states and store them in a memory. Later, we randomly choose a sample from the memory and calculate 2 different kinds of losses Loss_policy and Loss_value. We use Variable(…, requires_grad = True) on the calculated losses. Then the code is:</p>
<pre><code class="lang-auto">loss = Lp + alpha*Lv
loss.mean().backward()
optimizer.step()
</code></pre>
<p>But, after each ‘subepisode’, I run</p>
<pre><code class="lang-auto">for name, param in model.named_parameters():
     if param.requires_grad:
           print(name, param.data)
</code></pre>
<p>to print the trainable parameters. I find all the params that I want, but I am seeing that none of them is changing. That means, the model along with the rnn and policy is not updating. I don’t know whats going on.</p>

ptrblck: <aside class="quote no-group" data-username="Khabbab_Zakaria" data-post="1" data-topic="140647">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/khabbab_zakaria/48/33485_2.png" class="avatar"> Khabbab_Zakaria:</div>
<blockquote>
<p>We use Variable(…, requires_grad = True) on the calculated losses.</p>
</blockquote>
</aside>
<p>This sounds wrong as you are rewrapping the loss tensors and are thus detaching them from the computation graph besides <code>Variable</code>s being deprecated since PyTorch <code>0.4</code>, so remove this line of code and rerun your training.</p>

Khabbab_Zakaria: <p>Okay… Thanks, this sounds correct.<br>
So, I did that. The complete policy code is this:</p>
<pre><code class="lang-auto">class Policy(nn.Module):
    """
    implements both actor and critic in one model
    """
    def __init__(self):
        super(Policy, self).__init__()
        self.fc1 = nn.Linear(state_size+1, 128)

        self.fc2 = nn.Linear(128, 64)

        # actor's layer
        self.action_head = nn.Linear(64, action_size)
        self.mu = nn.Sigmoid()
        self.var = nn.Softplus()

        # critic's layer
        self.value_head = nn.Linear(64, 1)


    def forward(self, x):
        """
        forward of both actor and critic
        """
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))

        # actor: choses action to take from state s_t 
        # by returning probability of each action
        action_prob = self.action_head(x)
        mu = self.mu(action_prob)
        var = self.var(action_prob)

        # critic: evaluates being in the state s_t
        state_values = self.value_head(x)

        return mu, var, state_values
policy = Policy()
</code></pre>
<p>In model class, we are calling this policy after the rnn. You can see, this is actor critic. And in agent class’s act method, we are calling the model to get the action like this:</p>
<pre><code class="lang-auto">    def act(self, some_input, state):
      mu, var, state_value = self.model(some_input, state)
      mu = mu.data.cpu().numpy()
      sigma = torch.sqrt(var).data.cpu().numpy()
      action = np.random.normal(mu, sigma)
      action = np.clip(action, 0, 1)
      action = torch.from_numpy(action/1000)
      return action, state_value
</code></pre>
<p>I must mention that in optimizer, we are calling the model.parameters. When we print all the trainable parameters in each epoch, we see that everything else is changing except for the policy.action_head. Any idea why this is happening? I am not using Variable any more. I must also mention how the losses are calculated now:</p>
<pre><code class="lang-auto">       advantage = reward - Value
       Lp = -math.log(pdf_prob_now)*advantage
       policy_losses.append(Lp)
       #similar for value_losses
#after all the runs in the epoch is done
loss = torch.stack(policy_losses).sum() + alpha*torch.stack(value_losses).sum()
loss.backward()
</code></pre>
<p>Here Value is the state_value (the 2nd output from agent.act) and the pdf_prob_now is the probability of the action from all possible actions which is calculated like this:</p>
<pre><code class="lang-auto">def find_pdf(policy, action, rnn_output):
    mu, var, _ = policy(rnn_output)
    mu = mu.data.cpu().numpy()
    sigma = torch.sqrt(var).data.cpu().numpy()
    pdf_probability = stats.norm.pdf(action.cpu(), loc=mu, scale=sigma)
    return pdf_probability
</code></pre>

ptrblck: <p>I don’t know how and when each code snippet is called, but check the <code>.grad</code> attribute of <code>policy.action_head</code> before and after the first <code>backward</code> call and make sure its set to <code>None</code> first and shows a valid gradient afterwards (remove the <code>zero_grad</code> call if it was used before).<br>
If you are seeing valid gradient values (even if they are close to zero) the gradients are indeed calculated but might be too small to see a significant change and you should then debug the used operations.</p>