Fighting against distributions Categorical: log_prob is delivering unexpected values

username: wilhelm
cooked: <p>I’m trying to implement the CartPole-v1 problem from the Book “Introduction into RL” by trying to adapt their code posted in their repository here: <a href="https://github.com/DeepReinforcementLearning/DeepReinforcementLearningInAction/blob/master/Chapter%204/Ch4_book.ipynb" rel="noopener nofollow ugc">link</a>.</p>
<p>Their code is old, so I changed it for working with Gymnasium 0.28.1 and I removed this line of code:</p>
<blockquote>
<p>action = np.random.choice(np.array([0,1]), p=act_prob.data.numpy()) <span class="hashtag-raw">#E</span></p>
</blockquote>
<p>with the code here:</p>
<blockquote>
<p>policy_dist = Categorical(probs = act_prob)<br>
action = policy_dist.sample()<br>
log_prob = policy_dist.log_prob(action)<br>
action = action.cpu().numpy()</p>
</blockquote>
<p>because I want to sample action from the Categorical distribution instead of using numpy to do that.<br>
Here is the code, which should run once copied into your system.</p>
<blockquote>
<p>import numpy as np<br>
import torch<br>
import gymnasium as gyms</p>
<p>from torch.distributions import Categorical</p>
<p>env = gyms.make(“CartPole-v1”)</p>
<p>l1 = 4<br>
l2 = 150<br>
l3 = 2</p>
<p>model = torch.nn.Sequential(<br>
torch.nn.Linear(l1, l2),<br>
torch.nn.LeakyReLU(),<br>
torch.nn.Linear(l2, l3),<br>
torch.nn.Softmax(dim=0)<br>
)</p>
<p>learning_rate = 0.009<br>
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)</p>
<p>def discount_rewards(rewards, gamma=0.99):<br>
lenr = len(rewards)<br>
disc_return = torch.pow(gamma,torch.arange(lenr).float()) * rewards<br>
disc_return /= disc_return.max()</p>
<pre><code>return disc_return
</code></pre>
<p>def loss_fn(preds, r):<br>
return -1 * torch.sum(r * torch.log(preds))</p>
<p>MAX_DUR = 499<br>
MAX_EPISODES = 500<br>
gamma = 0.99<br>
scores = <span class="chcklst-box fa fa-square-o fa-fw"></span><br>
expectation = 0.0</p>
<p>for episode in range(MAX_EPISODES):<br>
curr_state, info = env.reset()<br>
termination = False<br>
truncation = False<br>
transitions = <span class="chcklst-box fa fa-square-o fa-fw"></span><br>
score = 0</p>
<pre><code>for t in range(MAX_DUR):
    act_prob = model(torch.from_numpy(curr_state).float())
    policy_dist = Categorical(probs = act_prob)
    action = policy_dist.sample()
    log_prob = policy_dist.log_prob(action)
    action = action.cpu().numpy()

    prev_state = curr_state
    curr_state, reward, termination, truncation, info = env.step(action)
    score += reward
    transitions.append((prev_state, action, log_prob, t+1))

    if termination or truncation:
        break

optimizer.zero_grad()

reward_batch = torch.Tensor([r for (s,a,lp, r) in transitions]).flip(dims=(0,))

disc_returns = discount_rewards(reward_batch)

state_batch = torch.stack([torch.from_numpy(s) for (s,a,lp,r) in transitions])
action_batch = torch.stack([torch.from_numpy(a) for (s,a,lp,r) in transitions])
log_prob = torch.stack([lp for (s,a,lp,r) in transitions])

pred_batch = model(state_batch)

prob_batch = pred_batch.gather(dim=1,index=action_batch.long().view(-1,1)).squeeze()

# Here the comparation
loss = loss_fn(prob_batch, disc_returns)

new_loss = -1 * torch.sum(disc_returns * log_prob)

loss.backward()
optimizer.step()

scores.append(score)
avg_score = np.mean(scores[-100:])

print(f"Episode: {episode + 1}, avg_score: {avg_score}")
</code></pre>
</blockquote>
<p>Now, I didn’t want to use the t<code>orch.log</code> call, because the Categorical distributions outputs the <code>log_prob</code> automatically. So I’m trying to calculate the loss differently. The original code does:</p>
<blockquote>
<p>loss = -1 * torch.sum(disc_returns * torch.log(probs))</p>
</blockquote>
<p>where probs is the output of the softmax function applied to the output of the network model.</p>
<p>I indeed I’m doing the following:</p>
<blockquote>
<p>new_loss = -1 * torch.sum(disc_returns * log_prob)</p>
</blockquote>
<p>where <code>log_prob</code> is given directly by the Categorical distribution  once the action is sampled:</p>
<blockquote>
<p>log_prob = policy_dist.log_prob(action)</p>
</blockquote>
<p>But I get not the same result.<br>
I outputted the result from the <code>torch.log(probs)</code> and from the <code>log_prob</code> operations and they are quite different:</p>
<blockquote>
<p>tensor([-2.6778, -2.6924, -2.7011, -2.6968, -2.7408, -2.6880, -2.6965, -2.7147, -2.6316, -2.6783, -2.6910, -2.7152, -2.6167, -2.6579, -2.6910],  grad_fn=)<br>
tensor([-0.6593, -0.6677, -0.6946, -0.6593, -0.6892, -0.6697, -0.6991, -0.7334, -0.6097, -0.6514, -0.7062, -0.7417, -0.6036, -0.6406, -0.7169],  grad_fn=)</p>
</blockquote>
<p>But they should be more or less the same.<br>
Any idea what’s wrong here?</p>

username: wilhelm
cooked: <p>Many thanks!!!</p>
<p>I have to admit, that I understand now more. But there are still some subtle problems, which are not completely clear to me.</p>
<p>By the way, here the full code with the changes:</p>
<blockquote>
<p>import numpy as np<br>
import torch<br>
import gymnasium as gyms</p>
<p>from torch.distributions import Categorical</p>
<p>torch.autograd.set_detect_anomaly(True)</p>
<p>env = gyms.make(“CartPole-v1”)</p>
<p>l1 = 4<br>
l2 = 150<br>
l3 = 2</p>
<p>model = torch.nn.Sequential(<br>
torch.nn.Linear(l1, l2),<br>
torch.nn.LeakyReLU(),<br>
torch.nn.Linear(l2, l3),<br>
torch.nn.Softmax(dim = 1)<br>
)</p>
<p>learning_rate = 0.009<br>
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)</p>
<p>def discount_rewards(rewards, gamma=0.99):<br>
lenr = len(rewards)<br>
disc_return = torch.pow(gamma,torch.arange(lenr).float()) * rewards<br>
disc_return /= disc_return.max()<br>
return disc_return</p>
<p>def loss_fn(preds, r):<br>
return -1 * torch.sum(r * torch.log(preds))</p>
<p>MAX_DUR = 499<br>
MAX_EPISODES = 500<br>
gamma = 0.99<br>
scores = <span class="chcklst-box fa fa-square-o fa-fw"></span><br>
expectation = 0.0</p>
<p>for episode in range(MAX_EPISODES):<br>
curr_state, info = env.reset()<br>
termination = False<br>
truncation = False<br>
transitions = <span class="chcklst-box fa fa-square-o fa-fw"></span><br>
score = 0</p>
<pre><code>for t in range(MAX_DUR):
    act_prob = model(torch.from_numpy(curr_state).unsqueeze(dim = 0).float())
    policy_dist = Categorical(probs = act_prob)
    action = policy_dist.sample()
    log_prob = policy_dist.log_prob(action)
    action = action.cpu().numpy()[0]

    prev_state = curr_state
    curr_state, reward, termination, truncation, info = env.step(action)
    score += reward
    transitions.append((prev_state, action, log_prob, t+1))

    if termination or truncation:
        break

ep_len = len(transitions)

optimizer.zero_grad()

reward_batch = torch.Tensor([r for (s,a,lp, r) in transitions]).flip(dims=(0,))

disc_returns = discount_rewards(reward_batch)

state_batch = torch.stack([torch.from_numpy(s) for (s,a,lp,r) in transitions])
action_batch = torch.tensor([a for (s,a,lp,r) in transitions])
log_prob = torch.stack([lp for (s,a,lp,r) in transitions])

pred_batch = model(state_batch)

prob_batch = pred_batch.gather(dim=1,index=action_batch.long().view(-1,1)).squeeze()

# Here the comparation
loss = loss_fn(prob_batch, disc_returns)

new_loss = -1 * torch.sum(disc_returns * log_prob.view(-1))

new_loss.backward()
optimizer.step()

scores.append(score)
avg_score = np.mean(scores[-100:])

print(f"Episode: {episode + 1}, avg_score: {avg_score:.4f}")
</code></pre>
</blockquote>
<p>If I print the two variables <code>torch.log(prob_batch)</code> and <code>log_prob</code>, then I get exact the same output:</p>
<blockquote>
<p>tensor([-0.7508, -0.6382, -0.7480, -0.6418, -0.7458, -0.7440, -0.7561, -0.7740, -0.7822, -0.6004, -0.7743, -0.7887, -0.8136, -0.5656, -0.8112], grad_fn= LogBackward0&gt;)<br>
tensor([-0.7508, -0.6382, -0.7480, -0.6418, -0.7458, -0.7440, -0.7561, -0.7740, -0.7822, -0.6004, -0.7743, -0.7887, -0.8136, -0.5656, -0.8112], grad_fn= ViewBackward0&gt;)</p>
</blockquote>
<p>If I calculate the loss using the <code>torch.log(prob_batch)</code>, then the algorithm learns in the 500 games<br>
But as far as I use the <code>log_prob</code>, then the same algorithm in the same conditions does not learns anything.<br>
I think it depends on the <code>grad_fn</code> variable. In one case is <code>grad_fn=&lt;LogBackward0&gt;</code>, in the other case is <code>grad_fn=&lt;ViewBackward0&gt;</code> . This is the only difference.<br>
What else could lead to the different behavior?</p>