How can I improve this machine learning algorithm so that the agent minimizes the location and angle in the observation of the cartpole environment?

```python

from torch import nn
import gymnasium as gym
import numpy as np
import torch
import torch.nn.functional as F
import torch.optim as optim


class NeuralNetwork(nn.Module):
    def __init__(self):
        super(NeuralNetwork, self).__init__()
        self.linear_stack = nn.Sequential(
            nn.Linear(4, 11),
            nn.LeakyReLU(),
            nn.Linear(11, 2),
        )

    def forward(self, x):
        x = self.linear_stack(x)
        return x


def train(model):
    optimizer = optim.AdamW(model.parameters(), lr=0.01)
    env = gym.make("CartPole-v1")

    for episode in range(300):  # Number of episodes to train on
        observation, info = env.reset()
        obs_tensor = torch.tensor(observation, dtype=torch.float32)

        log_probs = []
        rewards = []

        for step in range(500):  # Max steps per episode
            # Forward pass
            action_logits = model(obs_tensor)
            action_probs = F.softmax(action_logits, dim=0)

            # Sample action from the policy
            m = torch.distributions.Categorical(action_probs)
            action = m.sample()
            log_prob = m.log_prob(action)

            # Step in the environment
            observation, reward, terminated, truncated, info = env.step(action.item())
            obs_tensor = torch.tensor(observation, dtype=torch.float32)
            reward /= np.linalg.norm(observation * np.array([0.5, 0, 1, 0]))

            # Save log probability and reward
            log_probs.append(log_prob)
            rewards.append(reward)

            if terminated or truncated:
                break

        # Compute cumulative discounted rewards
        R = 0
        gamma = 0.99  # Discount factor
        discounted_rewards = []
        for r in rewards[::-1]:
            R = r + gamma * R
            discounted_rewards.insert(0, R)
        discounted_rewards = torch.tensor(discounted_rewards)

        # Normalize rewards
        discounted_rewards = (discounted_rewards - discounted_rewards.mean()) / (
            discounted_rewards.std() + 1e-9
        )

        # Compute loss
        policy_losses = []
        for log_prob, R in zip(log_probs, discounted_rewards):
            policy_losses.append(-log_prob * R)
        loss = torch.stack(policy_losses).sum()

        # Perform backpropagation
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        total_reward = sum(rewards)
        print(f"Episode {episode + 1}\tTotal Reward: {total_reward}")

    env.close()


def main():

    model = NeuralNetwork()
    train(model)

    env = gym.make("CartPole-v1", render_mode="human")

    observation, info = env.reset()

    for step in range(200):
        obs_tensor = torch.tensor(observation, dtype=torch.float32)
        with torch.no_grad():
            action_logits = model(obs_tensor)
            action = torch.argmax(action_logits)

        observation, reward, terminated, truncated, info = env.step(action.item())

        if terminated or truncated:
            observation, info = env.reset()

    env.close()


if __name__ == "__main__":
    main()

```