Unidentifiable flipped sign in policy gradient

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, hidd…