Deep RL problem: Loss decreases but agent doesn't learn

I'm implementing a basic Vanilla Policy Gradient algorithm for the CartPole-v1 gymnasium environment, and I don't know what I'm doing wrong. No matter what I try, during the training loop the loss decreases (so the model is actually learning something), but the episode total reward also decreases until it reaches around 9-10 steps (which I imagine is the minimum number of steps needed to make the pole fall). So it's learning to do it bad! In this algorithm there are a few moving parts that influence each other, and I run out of ideas to see where the problem could be. For the discounted rewards I'm using $ Q_{k,t} = \sum_{i=0}{\gamma^{i-t} r_i} $ for $k$ = all episodes and $t$ = all steps in an episode. And for the loss: $ L = -\sum_{k,t}Q_{k,t}log\pi_{\theta}(a_t | s_t)$ The code is a mix from Maxim Lapan's Deep RL Hands-On book, Karpathy's Pong example (blog, code), and personal tweaks. Here's my code: import gymnasium as gym import torch from torch import nn import torch.nn.functional as F from torch.nn.init import xavier_uniform_ import numpy as np GAMMA = 0.99 LEARNING_RATE = 0.001 BATCH_SIZE = 4 DEVICE = torch.device('mps') class XavierLinear(nn.Linear): def __init__(self, in_features: int, out_features: int, bias: bool = True, device=None, dtype=None) -> None: super().__init__(in_features, out_features, bias, device, dtype) xavier_uniform_(self.weight) class VPG(nn.Module): def __init__(self, input_size, output_size): super(VPG, self).__init__() self.net = nn.Sequenti…