Answers.Body: <p>Turns out the .unsqueeze(1) for the values tensor is what messed everything up.
  Still not quite sure how that happened.</p>
  
Answers.ContentLicense: CC BY-SA 4.0
Answers.CreationDate: 2024-11-27T18:23:50.473Z
Answers.LastActivityDate: 2024-11-27T18:23:50.473Z
Answers.UserDisplayName: Leonhard Piff
Body: <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>
  
Tags: <reinforcement-learning><deep-learning>
Title: Unidentifiable flipped sign in policy gradient
UserDisplayName: Leonhard Piff
ThreadText: # 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.
  
  
  ---
  
  