The gradients are non-zero when i use torch.autograd, but param.grad returns none and the step function does not update weights

gebrahimi: <p>I am trying to train the following policy gradient network, however the weights do not update. not sure what is the problem. Can someone help, please?</p>
<pre><code>class Policy(nn.Module):
def __init__(self, obs_shape, hidden_size = 64):
    super(Policy, self).__init__()
    
    init_ = lambda m: init(m, nn.init.orthogonal_, lambda x: nn.init.
                           constant_(x, 0), np.sqrt(2))

    self.actor = nn.Sequential(
        nn.Linear(obs_shape, hidden_size), nn.BatchNorm1d(hidden_size) , nn.ReLU(),
        nn.Linear(hidden_size, hidden_size),nn.BatchNorm1d(hidden_size), nn.ReLU()
        ,nn.Linear(hidden_size, 3))
    

def forward(self, inputs):
    x = inputs


    hidden_actor = self.actor(x)

    return hidden_actor


def act(self, inputs, deterministic=False):
    actor_features = self.forward(inputs)
    action_prob = F.softmax(actor_features, dim = -1)
            
    dist = distributions.Categorical(action_prob)
    if deterministic:
        action = dist.mode()
    else:
        action = dist.sample()

    action_log_probs = dist.log_prob(action)
    dist_entropy = dist.entropy().mean()

    return action, action_log_probs


def evaluate_actions(self, inputs, action):
    actor_features = self.forward(inputs)
    #dist = self.dist(actor_features)
    action_prob = F.softmax(actor_features, dim = -1)
    dist = distributions.Categorical(action_prob)
    action_log_probs = dist.log_prob(action)

    return action_log_probs
</code></pre>
<p>The train function is the following:</p>
<pre><code>def train(env, policy, optimizer, discount_factor):

policy.train()

log_prob_actions = []
rewards = []
states = []
actions = []
done = False
episode_reward = 0

state = env.reset()

while not done:

    state = torch.FloatTensor(state).unsqueeze(0).to(device)
    states.append(state)
    policy.eval()
    
    with torch.no_grad():
        action, log_prob_action = policy.act(state)

    policy.train()
    actions.append(action)
    

    
    state, reward, done, _ = env.step(action.item())
    #print(log_prob_action)
    log_prob_actions.append(log_prob_action)
    rewards.append(reward)

    episode_reward += reward

log_prob_actions = torch.cat(log_prob_actions)

returns = calculate_returns(rewards, discount_factor).to(device)
loss, grad_norm = update_policy_version2(returns, states, actions, optimizer)

return loss, grad_norm, episode_reward, states, actions, returns, rewards
</code></pre>
<p>Where update policy function is:</p>
<pre><code>def update_policy_version2(returns, states, actions, optimizer):
states = torch.cat(states)
actions = torch.stack(actions).squeeze()
log_prob_actions = policy.evaluate_actions(states, actions)

returns = returns.detach()

loss = - (returns * log_prob_actions).sum()

optimizer.zero_grad()
grad_loss = torch.autograd.grad(loss, policy.parameters(), retain_graph=True)
#print(grad_loss)
grad_loss = torch.cat([g.contiguous().view(-1) for g in grad_loss])
norm =  torch.norm(grad_loss)

optimizer.step()
params = []
for param in policy.parameters():
    params.append(param.view(-1))

for p in policy.parameters():
    print(p.grad)

return loss.item(), norm.item()
</code></pre>
<p>The print function for p.grad returns none while the grad_loss calculated using torch.autograd.grad returns real values. The step() function does not update the weights.</p>

ptrblck: <p><code>torch.autograd.grad</code> does not manipulate the <code>.grad</code> attributes, but returns the desired gradients as the return value.<br>
Here is a small example for resnet18, where I only calculate the gradients for the <code>weight</code> parameter of the first convolution layer.<br>
Note that even though I get the gradients as the return value, the <code>.grad</code> attribute is still <code>None</code>:</p>
<pre><code class="lang-python">model = models.resnet18()

out = model(torch.randn(1, 3, 224, 224))
grad = torch.autograd.grad(out.mean(), model.conv1.weight)
print(grad[0].shape == model.conv1.weight.shape)
&gt; True
for name, param in model.named_parameters():
    print(name, param.grad)
</code></pre>
<p>So you would have to apply the gradients manually or let Autograd calculate and assign them to the <code>.grad</code> attributes via <code>loss.backward()</code>.</p>
<p>I’m not sure if this might be related, but you are also detaching the <code>returns</code> tensor.<br>
I’m not familiar with your use case, so it might work as expected.</p>