How can I make this machine learning algorithm for continuous action learn better?

```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(3, 4),
            nn.LeakyReLU(),
            nn.Linear(4, 11),
            nn.LeakyReLU()
        )

        self.mean = nn.Linear(11, 1)
        self.std = nn.Linear(11, 1)

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

        mean = self.mean(x)
        std = self.std(x)

        std = F.sigmoid(std)

        return mean, std


def train(model):
    optimizer = optim.AdamW(model.parameters(), lr=0.01)
    env = gym.make("Pendulum-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
            mean, std = model(obs_tensor)
            normal_distributions = torch.distributions.normal.Normal(mean, std)
            action = normal_distributions.sample()
            log_prob = normal_distributions.log_prob(action)

            action_np = np.clip(action.detach().numpy(), -2, 2)

            # Step in the environment
            observation, reward, terminated, truncated, info = env.step(action_np)
            obs_tensor = torch.tensor(observation, dtype=torch.float32)

            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-8
        )

        # 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()
    model.load_state_dict(torch.load("pendulum_1.pt", weights_only=True))

    # train(model)
    # torch.save(model.state_dict(), "pendulum_1.pt")

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

    observation, info = env.reset()

    for step in range(300):
        obs_tensor = torch.tensor(observation, dtype=torch.float32)
        with torch.no_grad():
            mean, _ = model(obs_tensor)
            action = np.clip(mean.detach().numpy(), -2, 2)
            print(action)

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

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

    env.close()


if __name__ == "__main__":
    main()

```