Actúa como un programador experto en algoritmos de reinforcement learning. Ahora revisa exhaustivamente el siguiente código de python en el que se implementa un algoritmo Actor-Critic para entrenar un agente en el ambiente Pendulum de gym y muéstrame partes del código que debo cambiar, pues durante el entrenamiento no logro aumentar la recompensa promedio. Te en cuenta que esto ocurre sólo para este ambiente que tiene espacio de acciones continuo, ya que cuando uso un ambiente con acciones discretas como "CartPole" sí funciona bien el entrenamiento. Aquí está el código:

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


class Actor(nn.Module):

    def __init__(self, dim_states, dim_actions, continuous_control):
        super(Actor, 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(1, 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 Critic(nn.Module):

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

    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 ActorCriticAgent:

    def __init__(self, dim_states, dim_actions, actor_lr, critic_lr, gamma, continuous_control=False):
        
        self._actor_lr = actor_lr
        self._critic_lr = critic_lr
        self._gamma = gamma

        self._dim_states = dim_states
        self._dim_actions = dim_actions

        self._continuous_control = continuous_control

        self._actor = Actor(self._dim_states, self._dim_actions, self._continuous_control)

        # Adam optimizer
        self._actor_optimizer = torch.optim.Adam(self._actor.parameters(), lr=self._actor_lr)

        self._critic = Critic(self._dim_states)

        # Adam optimizer
        self._critic_optimizer = torch.optim.Adam(self._critic.parameters(), lr=self._critic_lr)

        self._select_action = self._select_action_continuous if self._continuous_control else self._select_action_discrete
        self._compute_actor_loss = self._compute_actor_loss_continuous if self._continuous_control else self._compute_actor_loss_discrete


    def select_action(self, observation):
        return self._select_action(observation)
        

    def _select_action_discrete(self, observation):
        # sample from categorical distribution
        with torch.no_grad():
            observation = torch.from_numpy(observation).float()
            logits = self._actor(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():
            observation = torch.from_numpy(observation).float()
            mean = self._actor(observation)
            std = self._actor._log_std.exp()
            normal = torch.distributions.Normal(mean, std)
            action = normal.sample()
        return action.numpy()[0]


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


    def _compute_actor_loss_continuous(self, observation_batch, action_batch, advantage_batch):
        # use negative logprobs * advantages
        mean = self._actor(observation_batch)
        std = self._actor._log_std.exp()
        normal = torch.distributions.Normal(mean, std)
        log_probs = normal.log_prob(action_batch).squeeze()

        actor_loss = -(log_probs * advantage_batch).mean()
        return actor_loss


    def _compute_critic_loss(self, observation_batch, reward_batch, next_observation_batch, done_batch):
        # minimize mean((r + gamma * V(s_t1) - V(s_t))^2)
        q_values = self._critic(observation_batch).squeeze()
        next_q_values = self._critic(next_observation_batch).squeeze()
        expected_q_values = reward_batch + self._gamma * next_q_values * (1 - done_batch)

        critic_loss = nn.functional.mse_loss(q_values, expected_q_values)

        return critic_loss


    def update_actor(self, observation_batch, action_batch, reward_batch, next_observation_batch, done_batch):
        # compute the advantages using the critic and update the actor parameters
        # use self._compute_actor_loss
        observation_tensor = torch.from_numpy(observation_batch).float()
        action_tensor = torch.from_numpy(action_batch).float()
        reward_tensor = torch.from_numpy(reward_batch).float()
        next_observation_tensor = torch.from_numpy(next_observation_batch).float()
        done_tensor = torch.from_numpy(done_batch).float()

        q_values = self._critic(observation_tensor).squeeze()
        next_q_values = self._critic(next_observation_tensor).squeeze()
        expected_q_values = reward_tensor + self._gamma * next_q_values * (1 - done_tensor)
        advantage_tensor = expected_q_values - q_values

        actor_loss = self._compute_actor_loss(observation_tensor, action_tensor, advantage_tensor)

        self._actor_optimizer.zero_grad()
        actor_loss.backward()
        self._actor_optimizer.step()
        
        
    def update_critic(self, observation_batch, reward_batch, next_observation_batch, done_batch):
        # update the critic
        # use self._compute_critic_loss
        observation_tensor = torch.from_numpy(observation_batch).float()
        reward_tensor = torch.from_numpy(reward_batch).float()
        next_observation_tensor = torch.from_numpy(next_observation_batch).float()
        done_tensor = torch.from_numpy(done_batch).float()

        critic_loss = self._compute_critic_loss(observation_tensor, reward_tensor, next_observation_tensor, done_tensor)
        self._critic_optimizer.zero_grad()
        critic_loss.backward()
        self._critic_optimizer.step()


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

    # Modify this function to return a tuple of numpy arrays containing:
    # (np.array(obs_t), np.array(acs_t), np.arraw(rew_t), np.array(obs_t1), np.array(done_t))
    # np.array(obs_t)   -> shape: (time_steps, nb_obs)
    # np.array(obs_t1)  -> shape: (time_steps, nb_obs)
    # np.array(acs_t)   -> shape: (time_steps, nb_acs) if actions are continuous, (time_steps,) if actions are discrete
    # np.array(rew_t)   -> shape: (time_steps,)
    # np.array(done_t)  -> shape: (time_steps,)

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

    obs_t_list, acs_t_list, rew_t_list, obs_t1_list, done_t_list = [], [], [], [], []

    while not done:

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

        action = agent.select_action(obs_t)
        
        obs_t1, reward, done, _ = env.step(action)

        obs_t = np.squeeze(obs_t)
        obs_t1 = np.squeeze(obs_t1)
        
        obs_t_list.append(obs_t)
        acs_t_list.append(action)
        rew_t_list.append(reward)
        obs_t1_list.append(obs_t1)
        done_t_list.append(done)

        obs_t = obs_t1
        episode_reward += reward
        
        nb_steps += 1

        if done:
            return np.array(obs_t_list), np.array(acs_t_list), np.array(rew_t_list), np.array(obs_t1_list), np.array(done_t_list)


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

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

        sampled_rollouts.append(sample_rollout)

    return sampled_rollouts


def train_agent(env, agent, training_iterations, min_batch_steps, nb_critic_updates):

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

    for tr_iter in range(training_iterations + 1):

        # Sample rollouts using sample_rollouts
        sampled_rollouts = sample_rollouts(env, agent, tr_iter, min_batch_steps)

        # performed_batch_steps >= min_batch_steps
        # Parse samples into the following arrays:

        sampled_obs_t  = np.concatenate([rollout[0] for rollout in sampled_rollouts])   # sampled_obs_t:  Numpy array, shape: (performed_batch_steps, nb_observations)
        sampled_acs_t  = np.concatenate([rollout[1] for rollout in sampled_rollouts])   # sampled_acs:    Numpy array, shape: (performed_batch_steps, nb_actions) if actions are continuous, 
                                                                                        #                                     (performed_batch_steps,)            if actions are discrete
        sampled_rew_t  = np.concatenate([rollout[2] for rollout in sampled_rollouts])   # sampled_rew_t:  Numpy array, shape: (performed_batch_steps,)
        sampled_obs_t1 = np.concatenate([rollout[3] for rollout in sampled_rollouts])   # sampled_obs_t1: Numpy array, shape: (performed_batch_steps, nb_observations)
        sampled_done_t = np.concatenate([rollout[4] for rollout in sampled_rollouts])   # sampled_done_t: Numpy array, shape: (performed_batch_steps,)

        # performance metrics
        update_performance_metrics(tr_iter, sampled_rollouts, axes, tr_iters_vec, avg_reward_vec, std_reward_vec, avg_steps_vec)

        for _ in range(nb_critic_updates):
            agent.update_critic(sampled_obs_t, sampled_rew_t, sampled_obs_t1, sampled_done_t)
        
        agent.update_actor(sampled_obs_t, sampled_acs_t, sampled_rew_t, sampled_obs_t1, sampled_done_t)

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

    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)

    # logs 
    print('-' * 32)
    print('%20s : %5d'   % ('Training iter'     ,(tr_iter)              ))
    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             ))

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

    avg_steps_vec.append(avg_steps)

    tr_iters_vec.append(tr_iter)

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


def plot_performance_metrics(axes, tr_iters_vec, avg_reward_vec, std_reward_vec, avg_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.plot(tr_iters_vec, avg_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

    actor_critic_agent = ActorCriticAgent(dim_states=dim_states,
                                          dim_actions=dim_actions,
                                          actor_lr=0.001,
                                          critic_lr=0.001,
                                          gamma=0.99,
                                          continuous_control=continuous_control)

    train_agent(env=env, 
                agent=actor_critic_agent, 
                training_iterations=2000, 
                min_batch_steps=5000,
                nb_critic_updates=1)