Actúa como un experto en algoritmos de reinforcement learning. Ahora revisa exhaustivamente el siguiente código de python en el que se implementa el entrenamiento del algoritmo REINFORCE en el ambiente Pendulum de gym y muéstrame secciones del código que se deban modificar sabiendo que durante el entrenamiento no se logra aumentar la recompensa promedio, ni tener un buen desempeño:

import gym
import time
import datetime
import csv
import torch 
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt

class Policy(nn.Module):

    def __init__(self, dim_states, dim_actions, continuous_control):
        super(Policy, self).__init__()
        # MLP, fully connected layers, ReLU activations, linear ouput activation
        # dim_states -> 64 -> 64 -> dim_actions
        self.layer1 = nn.Linear(dim_states, 64)
        self.layer2 = nn.Linear(64, 64)
        self.layer3 = nn.Linear(64, dim_actions)

        if continuous_control:
            # trainable parameter
            self._log_std = nn.Parameter(torch.zeros(dim_actions))


    def forward(self, input):
        input = nn.functional.relu(self.layer1(input))
        input = nn.functional.relu(self.layer2(input))
        input = self.layer3(input)
        return input


class PolicyGradients:

    def __init__(self, dim_states, dim_actions, lr, gamma, 
                 continuous_control=False, reward_to_go=False, use_baseline=False):
        
        self._learning_rate = lr
        self._gamma = gamma
        
        self._dim_states = dim_states
        self._dim_actions = dim_actions

        self._continuous_control = continuous_control
        self._use_reward_to_go = reward_to_go
        self._use_baseline = use_baseline

        self._policy = Policy(self._dim_states, self._dim_actions, self._continuous_control)
        # Adam optimizer
        self._optimizer = torch.optim.Adam(self._policy.parameters(), lr=self._learning_rate)

        self._select_action = self._select_action_continuous if self._continuous_control else self._select_action_discrete
        self._compute_loss = self._compute_loss_continuous if self._continuous_control else self._compute_loss_discrete


    def select_action(self, obs):
        return self._select_action(torch.tensor(obs, dtype=torch.float32))
        

    def _select_action_discrete(self, observation):
        # sample from categorical distribution
        with torch.no_grad():
            logits = self._policy(observation)
            dist = torch.distributions.Categorical(logits=logits)
            action = dist.sample()
        return action.item()


    def _select_action_continuous(self, observation):
        # sample from normal distribution
        # use the log std trainable parameter
        with torch.no_grad():
            mu = self._policy(observation)
            std = torch.exp(self._policy._log_std)
            dist = torch.distributions.Normal(mu, std)
            action = dist.sample()
        return action.numpy()
            

    def update(self, observation_batch, action_batch, advantage_batch):
        # update the policy here
        # you should use self._compute_loss
        observation_batch = torch.tensor(observation_batch, dtype=torch.float32)
        action_batch = torch.tensor(action_batch, dtype=torch.float32)
        advantage_batch = torch.tensor(advantage_batch, dtype=torch.float32)
        loss = self._compute_loss(observation_batch, action_batch, advantage_batch)
        self._optimizer.zero_grad()
        loss.backward()
        self._optimizer.step()
    

    def _compute_loss_discrete(self, observation_batch, action_batch, advantage_batch):
        # use negative logprobs * advantages
        logits = self._policy(observation_batch)
        dist = torch.distributions.Categorical(logits=logits)
        log_prob = dist.log_prob(action_batch)
        return -(log_prob * advantage_batch).mean()


    def _compute_loss_continuous(self, observation_batch, action_batch, advantage_batch):
        # use negative logprobs * advantages
        mu = self._policy(observation_batch)
        std = torch.exp(self._policy._log_std)
        dist = torch.distributions.Normal(mu, std)
        log_prob = dist.log_prob(action_batch).squeeze()
        return -(log_prob * advantage_batch).mean()

    
    def estimate_returns(self, rollouts_rew):
        estimated_returns = []
        for rollout_rew in rollouts_rew:
            discounted_rewards = self._discount_rewards(rollout_rew)

            if self._use_reward_to_go:
                # only for part 2
                n = len(rollout_rew)
                estimated_return = np.zeros_like(rollout_rew)
                for i in reversed(range(n)):
                    estimated_return[i] = rollout_rew[i] + (estimated_return[i+1] if i+1 < n else 0)
                estimated_return = list(estimated_return)

            else:
                total_reward = np.sum(discounted_rewards)
                estimated_return = [total_reward] * len(rollout_rew)

            estimated_returns += estimated_return

        if self._use_baseline:
            # only for part 2
            average_return_baseline = np.mean(estimated_returns)
            # Use the baseline:
            estimated_returns -= average_return_baseline

        # Normalize returns
        eps = np.finfo(np.float32).eps.item()
        estimated_returns = np.array(estimated_returns, dtype=np.float32)
        estimated_returns = (estimated_returns - np.mean(estimated_returns)) / (np.std(estimated_returns) + eps)
        return estimated_returns


    # It may be useful to discount the rewards using an auxiliary function [optional]
    def _discount_rewards(self, rewards):
        discounted_rewards = []
        discounted_reward = 0
        for reward in reversed(rewards):
            discounted_reward = reward + (discounted_reward * self._gamma)
            discounted_rewards.append(discounted_reward)
        discounted_rewards.reverse()

        return discounted_rewards


def perform_single_rollout(env, agent, episode_nb, render=False):

    # Modify this function to return a tuple of numpy arrays containing (observations, actions, rewards).
    # (np.array(obs), np.array(acs), np.array(rws))
    # np.array(obs) -> shape: (time_steps, nb_obs)
    # np.array(acs) -> shape: (time_steps, nb_acs) if actions are continuous, (time_steps,) if actions are discrete
    # np.array(rws) -> shape: (time_steps,)

    ob_t = env.reset()
    
    done = False
    episode_reward = 0
    nb_steps = 0

    obs = []
    acs = []
    rws = []

    while not done:

        if render:
            env.render()
            time.sleep(1. / 60)

        action = agent.select_action(ob_t)

        ob_t1, reward, done, _ = env.step(action)

        ob_t = np.squeeze(ob_t1) # <-- may not be needed depending on gym version
        
        obs.append(ob_t)
        acs.append(action)
        rws.append(reward)

        episode_reward += reward
        
        nb_steps += 1

        if done:
            return (np.array(obs), np.array(acs), np.array(rws))


def sample_rollouts(env, agent, training_iter, min_batch_steps):

    sampled_rollouts = []
    total_nb_steps = 0
    episode_nb = 0
    
    while total_nb_steps < min_batch_steps:

        episode_nb += 1
        render = 2 == 0 and len(sampled_rollouts) == 0 # Change training_iter%10 to any number you want

        # Use perform_single_rollout to get data 
        # Uncomment once perform_single_rollout works.
        # Return sampled_rollouts
        
        sample_rollout = perform_single_rollout(env, agent, episode_nb, render=render)
        total_nb_steps += len(sample_rollout[0])

        sampled_rollouts.append(sample_rollout)

        # if episode_nb == 2:
        #     break

    return sampled_rollouts


def train_pg_agent(env, agent, training_iterations, min_batch_steps):

    tr_iters_vec, avg_reward_vec, std_reward_vec, avg_steps_vec, std_steps_vec = [], [], [], [], []
    _, (axes) = plt.subplots(1, 2, figsize=(12,4))

    
    for tr_iter in range(training_iterations):
        # Sample rollouts using sample_rollouts
        sampled_rollouts = sample_rollouts(env, agent, tr_iter, min_batch_steps)

        # Parse sampled observations, actions and reward into three arrays:
        # performed_batch_steps >= min_batch_steps
        # sampled_obs: Numpy array, shape: (performed_batch_steps, dim_observations)
        sampled_obs = np.concatenate([rollout[0] for rollout in sampled_rollouts])
        # sampled_acs: Numpy array, shape: (performed_batch_steps, dim_actions) if actions are continuous, (performed_batch_steps,) if actions are discrete
        sampled_acs = np.concatenate([rollout[1] for rollout in sampled_rollouts])
        # sampled_rew: standard array of length equal to the number of trayectories that were sampled.
        # You may change the shape of sampled_rew, but it is useful keeping it as is to estimate returns.
        sampled_rew = [np.array(rollout[2]) for rollout in sampled_rollouts]

        # Return estimation
        # estimated_returns: Numpy array, shape: (performed_batch_steps, )
        estimated_returns = agent.estimate_returns(sampled_rew)
        
        # performance metrics
        update_performance_metrics(tr_iter, sampled_rollouts, axes, tr_iters_vec, avg_reward_vec, std_reward_vec, avg_steps_vec, std_steps_vec)

        agent.update(sampled_obs, sampled_acs, estimated_returns)
    
    save_metrics(tr_iters_vec,avg_reward_vec, std_reward_vec)
    

def update_performance_metrics(tr_iter, sampled_rollouts, axes, tr_iters_vec, avg_reward_vec, std_reward_vec, avg_steps_vec, std_steps_vec):

    raw_returns     = np.array([np.sum(rollout[2]) for rollout in sampled_rollouts])
    rollout_steps   = np.array([len(rollout[2]) for rollout in sampled_rollouts])

    avg_return = np.average(raw_returns)
    max_episode_return = np.max(raw_returns)
    min_episode_return = np.min(raw_returns)
    std_return = np.std(raw_returns)
    avg_steps = np.average(rollout_steps)
    std_steps = np.std(rollout_steps)

    # logs 
    print('-' * 32)
    print('%20s : %5d'   % ('Training iter'     ,(tr_iter + 1)          ))
    print('-' * 32)
    print('%20s : %5.3g' % ('Max episode return', max_episode_return    ))
    print('%20s : %5.3g' % ('Min episode return', min_episode_return    ))
    print('%20s : %5.3g' % ('Return avg'        , avg_return            ))
    print('%20s : %5.3g' % ('Return std'        , std_return            ))
    print('%20s : %5.3g' % ('Steps avg'         , avg_steps             ))
    print('%20s : %5.3g' % ('Steps std'         , std_steps             ))

    avg_reward_vec.append(avg_return)
    std_reward_vec.append(std_return)

    avg_steps_vec.append(avg_steps)
    std_steps_vec.append(std_steps)

    tr_iters_vec.append(tr_iter)

    plot_performance_metrics(axes, 
                            tr_iters_vec, 
                            avg_reward_vec, 
                            std_reward_vec, 
                            avg_steps_vec,
                            std_steps_vec)


def plot_performance_metrics(axes, tr_iters_vec, avg_reward_vec, std_reward_vec, avg_steps_vec, std_steps_vec):
    ax1, ax2 = axes
    
    [ax.cla() for ax in axes]
    ax1.errorbar(tr_iters_vec, avg_reward_vec, yerr=std_reward_vec, marker='.',color='C0')
    ax1.set_ylabel('Avg Reward')
    ax2.errorbar(tr_iters_vec, avg_steps_vec, yerr=std_steps_vec, marker='.',color='C1')
    ax2.set_ylabel('Avg Steps')

    [ax.grid('on') for ax in axes]
    [ax.set_xlabel('training iteration') for ax in axes]
    plt.pause(0.05)


def save_metrics(tr_iters_vec, avg_reward_vec, std_reward_vec):
    with open('metrics'+datetime.datetime.now().strftime('%H_%M_%S')+'.csv', 'w') as csv_file:
        csv_writer = csv.writer(csv_file, delimiter='\t')
        csv_writer.writerow(['steps', 'avg_reward', 'std_reward'])
        for i in range(len(tr_iters_vec)):
            csv_writer.writerow([tr_iters_vec[i], avg_reward_vec[i], std_reward_vec[i]])


if __name__ == '__main__':

    env = gym.make('Pendulum-v1')

    dim_states = env.observation_space.shape[0]

    continuous_control = isinstance(env.action_space, gym.spaces.Box)
    dim_actions = env.action_space.shape[0] if continuous_control else env.action_space.n

    policy_gradients_agent = PolicyGradients(dim_states=dim_states, 
                                             dim_actions=dim_actions, 
                                             lr=0.005,
                                             gamma=0.99,
                                             continuous_control=continuous_control,
                                             reward_to_go=False,
                                             use_baseline=False)
    
    train_pg_agent(env=env, 
                   agent=policy_gradients_agent, 
                   training_iterations=1000,
                   min_batch_steps=5000)