<p>Turns out the .unsqueeze(1) for the values tensor is what messed everything up.
Still not quite sure how that happened.</p>

CC BY-SA 4.0
2024-11-27T18:23:50.473Z
2024-11-27T18:23:50.473Z
Leonhard Piff
<p>Today I was building a VPG agent for a test and noticed it was getting worse not better over time so I flipped the reward during the training loop and lo and behold it learned. so obviously I started to look for where I flipped the sign. Problem is i've been looking for about 2 and a half hours now at a piece of code that took me 15 minutes to write and I can't figure it out:</p>
<pre><code>!pip install gymnasium
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.distributions import Categorical
import torch.optim as optimizer
import gymnasium as gym
import numpy as np
import random
import matplotlib.pyplot as plt
from IPython.display import clear_output

torch.set_default_device(&quot;cuda&quot; if torch.cuda.is_available() else &quot;cpu&quot;)

class FFN(nn.Module):
  def __init__(self, hidden_dim):
    super(FFN, self).__init__()
    self.norm = nn.LazyBatchNorm1d()
    self.fc1 = nn.LazyLinear(hidden_dim*4, bias=True)
    self.fc2 = nn.LazyLinear(hidden_dim, bias=True)
    self.activation = nn.SiLU()
  
  def forward(self, x):
    x2 = self.norm(x)
    x2 = self.fc1(x2)
    x2 = self.activation(x2)
    x2 = self.fc2(x2)
    return x + x2

class FFN(nn.Module):
  def __init__(self, input_dim):
    super(FFN, self).__init__()
    self.norm = nn.LayerNorm(input_dim)
    self.fc1 = nn.Linear(input_dim, input_dim*4)
    self.fc2 = nn.Linear(input_dim*4, input_dim)
    self.activation = nn.Mish()

  def forward(self, x):
    x2 = self.norm(x)
    x2 = self.fc1(x2)
    x2 = self.activation(x2)
    x2 = self.fc2(x2)
    return x + x2

class VPG(nn.Module):
  def __init__(self, hidden_dim=128, output_size=2, num_layers=2, clipnorm=0.001):
    super(VPG, self).__init__()
    self.clipnorm = clipnorm
    self.initial_layer = nn.LazyLinear(hidden_dim, bias=True)
    self.ffn = nn.ModuleList([FFN(hidden_dim) for _ in range(num_layers)])
    self.output_layer = nn.LazyLinear(output_size, bias=True)

    self.state_buffer = []
    self.action_buffer = []
    self.value_buffer = []

    self.episode_states = []
    self.episode_actions = []
    self.episode_rewards = []

    self.optimizer = optimizer.Adam(self.parameters(), lr=3e-4)

  def forward(self, x):
    x = self.initial_layer(x)
    for layer in self.ffn:
      x = layer(x)
    x = self.output_layer(x)
    dist = Categorical(logits = x)
    return dist

  def inference_forward(self, x):
    self.eval()
    with torch.no_grad():
      x = self.initial_layer(x)
      for layer in self.ffn:
        x = layer(x)
      x = self.output_layer(x)
      dist = Categorical(logits = x)
    self.train()
    return dist

  def get_action(self, state):
    state = torch.from_numpy(state).float()
    state = state.unsqueeze(0)
    action_dist = self.inference_forward(state)
    action = action_dist.sample()
    #print(action)
    return action.item()
  
  def add_experience(self, state, action, reward):
    self.episode_states.append(state)
    self.episode_actions.append(action)
    self.episode_rewards.append(reward)
  
  def integrate_episode(self):
    self.state_buffer.extend(self.episode_states)
    self.action_buffer.extend(self.episode_actions)
    values = discount(torch.tensor(self.episode_rewards))
    self.value_buffer.extend(values)

    self.episode_states = []
    self.episode_actions = []
    self.episode_rewards = []


  
  def train_step(self):

    state_tensor = torch.tensor(self.state_buffer, dtype=torch.float32)
    action_tensor = torch.tensor(self.action_buffer, dtype=torch.int32).unsqueeze(1)
    value_tensor = torch.tensor(self.value_buffer, dtype=torch.float32).unsqueeze(1)

    #print(f&quot;State Tensor: {state_tensor.shape}&quot;)
    #print(f&quot;Action Tensor: {action_tensor}&quot;)
    #print(f&quot;Value Tensor: {value_tensor.shape}&quot;


    action_dist = self.forward(state_tensor)
    log_probs = action_dist.log_prob(action_tensor)

    value_mean = value_tensor.mean()
    value_std = value_tensor.std()
    values = (value_tensor - value_mean) / value_std


    #print(f&quot;Values: {values}&quot;)
    loss = ((-log_probs * values.detach()).mean()) + (-0.001 * (action_dist.entropy().mean()))
  
    self.optimizer.zero_grad()
    loss.backward()
    torch.nn.utils.clip_grad_norm_(self.parameters(), self.clipnorm)
    self.optimizer.step()

    self.state_buffer = []
    self.action_buffer = []
    self.value_buffer = []






base_agent = VPG(hidden_dim=128, output_size=2, num_layers=1)

env = gym.make(&quot;CartPole-v1&quot;)

num_episodes_per_update = 4

reward_history = []

for episode in range(1000):
  state, info = env.reset()
  done = False
  total_reward = 0
  while not done:
    action = base_agent.get_action(state)
    new_state, reward, done, truncated, info = env.step(action)
    total_reward += reward
    base_agent.add_experience(state, action, -reward)
    state = new_state
  reward_history.append(total_reward)
  base_agent.integrate_episode()
  if episode % 10 == 0:
    print(f&quot;Episode: {episode} - Reward: {total_reward}&quot;)
    plt.plot(reward_history)
    plt.show()
    clear_output(wait=True)
  if episode % num_episodes_per_update == 0:
    base_agent.train_step()



<span class="math-container">```</span>
</code></pre>

<reinforcement-learning><deep-learning>
Unidentifiable flipped sign in policy gradient
Leonhard Piff
# Unidentifiable flipped sign in policy gradient

(Asked by: username_0 on 2024-11-27)

Today I was building a VPG agent for a test and noticed it was getting worse not better over time so I flipped the reward during the training loop and lo and behold it learned. so obviously I started to look for where I flipped the sign. Problem is i've been looking for about 2 and a half hours now at a piece of code that took me 15 minutes to write and I can't figure it out:
!pip install gymnasium
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.distributions import Categorical
import torch.optim as optimizer
import gymnasium as gym
import numpy as np
import random
import matplotlib.pyplot as plt
from IPython.display import clear_output

torch.set_default_device("cuda" if torch.cuda.is_available() else "cpu")

class FFN(nn.Module):
  def __init__(self, hidden_dim):
    super(FFN, self).__init__()
    self.norm = nn.LazyBatchNorm1d()
    self.fc1 = nn.LazyLinear(hidden_dim*4, bias=True)
    self.fc2 = nn.LazyLinear(hidden_dim, bias=True)
    self.activation = nn.SiLU()
  
  def forward(self, x):
    x2 = self.norm(x)
    x2 = self.fc1(x2)
    x2 = self.activation(x2)
    x2 = self.fc2(x2)
    return x + x2

class FFN(nn.Module):
  def __init__(self, input_dim):
    super(FFN, self).__init__()
    self.norm = nn.LayerNorm(input_dim)
    self.fc1 = nn.Linear(input_dim, input_dim*4)
    self.fc2 = nn.Linear(input_dim*4, input_dim)
    self.activation = nn.Mish()

  def forward(self, x):
    x2 = self.norm(x)
    x2 = self.fc1(x2)
    x2 = self.activation(x2)
    x2 = self.fc2(x2)
    return x + x2

class VPG(nn.Module):
  def __init__(self, hidden_dim=128, output_size=2, num_layers=2, clipnorm=0.001):
    super(VPG, self).__init__()
    self.clipnorm = clipnorm
    self.initial_layer = nn.LazyLinear(hidden_dim, bias=True)
    self.ffn = nn.ModuleList([FFN(hidden_dim) for _ in range(num_layers)])
    self.output_layer = nn.LazyLinear(output_size, bias=True)

    self.state_buffer = []
    self.action_buffer = []
    self.value_buffer = []

    self.episode_states = []
    self.episode_actions = []
    self.episode_rewards = []

    self.optimizer = optimizer.Adam(self.parameters(), lr=3e-4)

  def forward(self, x):
    x = self.initial_layer(x)
    for layer in self.ffn:
      x = layer(x)
    x = self.output_layer(x)
    dist = Categorical(logits = x)
    return dist

  def inference_forward(self, x):
    self.eval()
    with torch.no_grad():
      x = self.initial_layer(x)
      for layer in self.ffn:
        x = layer(x)
      x = self.output_layer(x)
      dist = Categorical(logits = x)
    self.train()
    return dist

  def get_action(self, state):
    state = torch.from_numpy(state).float()
    state = state.unsqueeze(0)
    action_dist = self.inference_forward(state)
    action = action_dist.sample()
    #print(action)
    return action.item()
  
  def add_experience(self, state, action, reward):
    self.episode_states.append(state)
    self.episode_actions.append(action)
    self.episode_rewards.append(reward)
  
  def integrate_episode(self):
    self.state_buffer.extend(self.episode_states)
    self.action_buffer.extend(self.episode_actions)
    values = discount(torch.tensor(self.episode_rewards))
    self.value_buffer.extend(values)

    self.episode_states = []
    self.episode_actions = []
    self.episode_rewards = []


  
  def train_step(self):

    state_tensor = torch.tensor(self.state_buffer, dtype=torch.float32)
    action_tensor = torch.tensor(self.action_buffer, dtype=torch.int32).unsqueeze(1)
    value_tensor = torch.tensor(self.value_buffer, dtype=torch.float32).unsqueeze(1)

    #print(f"State Tensor: {state_tensor.shape}")
    #print(f"Action Tensor: {action_tensor}")
    #print(f"Value Tensor: {value_tensor.shape}"


    action_dist = self.forward(state_tensor)
    log_probs = action_dist.log_prob(action_tensor)

    value_mean = value_tensor.mean()
    value_std = value_tensor.std()
    values = (value_tensor - value_mean) / value_std


    #print(f"Values: {values}")
    loss = ((-log_probs * values.detach()).mean()) + (-0.001 * (action_dist.entropy().mean()))
  
    self.optimizer.zero_grad()
    loss.backward()
    torch.nn.utils.clip_grad_norm_(self.parameters(), self.clipnorm)
    self.optimizer.step()

    self.state_buffer = []
    self.action_buffer = []
    self.value_buffer = []






base_agent = VPG(hidden_dim=128, output_size=2, num_layers=1)

env = gym.make("CartPole-v1")

num_episodes_per_update = 4

reward_history = []

for episode in range(1000):
  state, info = env.reset()
  done = False
  total_reward = 0
  while not done:
    action = base_agent.get_action(state)
    new_state, reward, done, truncated, info = env.step(action)
    total_reward += reward
    base_agent.add_experience(state, action, -reward)
    state = new_state
  reward_history.append(total_reward)
  base_agent.integrate_episode()
  if episode % 10 == 0:
    print(f"Episode: {episode} - Reward: {total_reward}")
    plt.plot(reward_history)
    plt.show()
    clear_output(wait=True)
  if episode % num_episodes_per_update == 0:
    base_agent.train_step()



```

**Tags:** reinforcement-learning, deep-learning


---

## 1 answers

### Answer by username_0 on 2024-11-27. Score: 0

Turns out the .unsqueeze(1) for the values tensor is what messed everything up.
Still not quite sure how that happened.


---

