Optimización de Políticas Proximales con PyTorch y Gymnasium

1. Configuración del entorno 

Antes de implementar PPO, necesitamos instalar las bibliotecas de software necesarias y elegir un entorno adecuado para aplicar la política. 

Instalación de PyTorch y las bibliotecas requeridas

Necesitamos instalar el siguiente software: 

• PyTorch y otras bibliotecas de software, como numpy (para funciones matemáticas y estadísticas) y matplotlib (para trazar gráficos).
• El paquete de software Gym de código abierto de OpenAI, una biblioteca de Python que simula diferentes entornos y juegos, que pueden resolverse utilizando el aprendizaje por refuerzo. Puedes utilizar la API de Gym para que tu algoritmo interactúe con el entorno. Dado que la funcionalidad de gym a veces cambia a través del proceso de actualización, en este ejemplo, congelamos su versión en 0.25.2.

Para instalar en un servidor o máquina local, ejecuta:

$ pip install torch numpy matplotlib gym==0.25.2

Para instalar usando un Notebook como Google Colab o DataLab, usa:

!pip install torch numpy matplotlib gym==0.25.2

Crear el(los) entorno(s) CartPole

Usa OpenAI Gym para crear dos instancias (una para entrenamiento y otra para pruebas) del entorno CartPole:

env_train = gym.make('CartPole-v1')
env_test = gym.make('CartPole-v1')

2. Implementando PPO en PyTorch

Ahora, implementemos PPO usando PyTorch.

Definiendo la red de política

Como se explicó anteriormente, PPO se implementa como un modelo actor-crítico. El actor implementa la política, y el crítico predice su valor estimado. Tanto las redes neuronales del actor como del crítico toman la misma entrada: el estado en cada paso de tiempo. Así, los modelos de actor y crítico pueden compartir una red neuronal común, que se denomina arquitectura de backbone. El actor y el crítico pueden extender la arquitectura de backbone con capas adicionales.

Define la red de backbone

Los siguientes pasos describen la red de backbone:

• Implementa una red con 3 capas: una de entrada, una oculta y una de salida.
• Después de las capas de entrada y oculta, usamos una función de activación. En este tutorial, elegimos ReLU porque es computacionalmente eficiente.
• También imponemos una función de abandono después de las capas de entrada y ocultas para obtener una red robusta. La función de abandono pone a cero aleatoriamente algunas neuronas. Esto reduce la dependencia de neuronas específicas y previene el sobreajuste, haciendo así que la red sea más robusta.

El código a continuación implementa la estructura básica:

class BackboneNetwork(nn.Module):
    def __init__(self, in_features, hidden_dimensions, out_features, dropout):
        super().__init__()
        self.layer1 = nn.Linear(in_features, hidden_dimensions)
        self.layer2 = nn.Linear(hidden_dimensions, hidden_dimensions)
        self.layer3 = nn.Linear(hidden_dimensions, out_features)
        self.dropout = nn.Dropout(dropout)
    def forward(self, x):
        x = self.layer1(x)
        x = f.relu(x)
        x = self.dropout(x)
        x = self.layer2(x)
        x = f.relu(x)
        x = self.dropout(x)
        x = self.layer3(x)
        return x

Definir la red actor-critic

Ahora, podemos usar esta red para definir la clase actor-critic, ActorCritic. El actor modela la política y predice la acción. El crítico modela la función de valor y predice el valor. Ambos toman el estado como entrada.

class ActorCritic(nn.Module):
    def __init__(self, actor, critic):
        super().__init__()
        self.actor = actor
        self.critic = critic
    def forward(self, state):
        action_pred = self.actor(state)
        value_pred = self.critic(state)
        return action_pred, value_pred

Instanciar las redes actor y crítico

Usaremos las redes definidas anteriormente para crear un actor y un crítico. Luego, crearemos un agente, incluyendo el actor y el crítico.

Antes de crear el agente, inicializa los parámetros de la red:

• Las dimensiones de la capa oculta, H, que es un parámetro configurable. El tamaño y número de capas ocultas dependen de la complejidad del problema. Usaremos una capa oculta con dimensiones 64 x 64.
• Características de entrada, N, donde N es el tamaño del array de estado. La capa de entrada tiene dimensiones N x H. En el entorno CartPole, el estado es un array de 4 elementos. Entonces, N es 4.
• Características de salida de la red actor, O, donde O es el número de acciones en el entorno. La capa de salida del actor tiene dimensiones H x O. El entorno CartPole tiene 2 acciones.
• Características de salida de la red crítica. Dado que la red crítica predice solo el valor esperado (dado un estado de entrada), el número de características de salida es 1.
• La deserción como una fracción.

El siguiente código muestra cómo declarar las redes de actor y crítico basadas en la red principal:

def create_agent(hidden_dimensions, dropout):
    INPUT_FEATURES = env_train.observation_space.shape[0]
    HIDDEN_DIMENSIONS = hidden_dimensions
    ACTOR_OUTPUT_FEATURES = env_train.action_space.n
    CRITIC_OUTPUT_FEATURES = 1
    DROPOUT = dropout
    actor = BackboneNetwork(
            INPUT_FEATURES, HIDDEN_DIMENSIONS, ACTOR_OUTPUT_FEATURES, DROPOUT)
    critic = BackboneNetwork(
            INPUT_FEATURES, HIDDEN_DIMENSIONS, CRITIC_OUTPUT_FEATURES, DROPOUT)
    agent = ActorCritic(actor, critic)
    return agent

Calculando los retornos

El entorno otorga una recompensa que va desde cada paso al siguiente, dependiendo de la acción del agente. La recompensa, R, se expresa como:

El retorno se define como el valor acumulado de las recompensas futuras esperadas. Las recompensas de pasos que están más lejos en el futuro son menos valiosas que las recompensas inmediatas. Por lo tanto, el retorno se calcula comúnmente como el retorno descontado, G, definido como:

En este tutorial (y muchas otras referencias), el retorno se refiere al retorno descontado. 

Para calcular el retorno:

• Comience con las recompensas esperadas de todos los estados futuros.
• Multiplica cada recompensa futura por un exponente del factor de descuento, . Por ejemplo, la recompensa esperada después de 2 pasos (desde el presente) se multiplica por 2. 
• Suma todas las recompensas futuras descontadas para calcular el retorno. 
• Normalice el valor del retorno. 

La función calcular_retornos() realiza estos cálculos, como se muestra a continuación:

def calculate_returns(rewards, discount_factor):
    returns = []
    cumulative_reward = 0
    for r in reversed(rewards):
        cumulative_reward = r + cumulative_reward * discount_factor
        returns.insert(0, cumulative_reward)
    returns = torch.tensor(returns)
    # normalizar el retorno
    returns = (returns - returns.mean()) / returns.std()
    return returns

Implementando la función ventaja

La ventaja se calcula como la diferencia entre el valor predicho por el crítico y el retorno esperado de las acciones elegidas por el actor de acuerdo con la política. Para una acción dada, la ventaja expresa el beneficio de tomar esa acción específica sobre una acción arbitraria (promedio).

En el documento PPO original (ecuación 10), la ventaja, mirando hacia adelante hasta el paso de tiempo T se expresa así:

Mientras se codifica el algoritmo, la restricción de mirar hacia adelante hasta un número establecido de pasos de tiempo se impone a través del tamaño del lote. Por lo tanto, la ecuación anterior se puede simplificar como la diferencia entre el valor y los retornos esperados. Los retornos esperados se cuantifican en la función de valor estado-acción, Q. 

Así, la fórmula simplificada a continuación expresa la ventaja de elegir: 

• una acción en particular 
• en un estado dado 
• bajo una política particular 
• en un paso de tiempo particular 

Esto se expresa como: 

OpenAI también utiliza esta fórmula para implementar RL. La función calculate_advantages() mostrada a continuación calcula la ventaja:

def calculate_advantages(returns, values):
    advantages = returns - values
    # Normalizar la ventaja
    advantages = (advantages - advantages.mean()) / advantages.std()
    return advantages

Pérdida de sustitución y mecanismo de recorte

La pérdida de política sería la pérdida estándar del gradiente de política sin técnicas especiales como PPO. La pérdida estándar del gradiente de política se calcula como el producto de:

• Las probabilidades de acción de la política
• La función de ventaja, que se calcula como la diferencia entre:
• El retorno de la política
• El valor esperado

La pérdida estándar del gradiente de política no puede corregir cambios abruptos en la política. La pérdida sustituta modifica la pérdida estándar para restringir la cantidad que la política puede cambiar en cada iteración. Es el mínimo de dos cantidades:

• El producto de:
• La proporción de la política. Esta proporción expresa la diferencia entre las probabilidades de acción antiguas y nuevas.
• La función de ventaja
• El producto de:
• El valor limitado de la proporción de la política. Esta proporción se recorta de manera que la política actualizada esté dentro de un cierto porcentaje de la política anterior.
• La función de ventaja

Para el proceso de optimización, la pérdida sustituta se utiliza como un proxy para la pérdida real.

El mecanismo de recorte

La razón de la política, R, es la diferencia entre las políticas nuevas y antiguas y se da como la razón de los logaritmos de las probabilidades de la política bajo los parámetros nuevos y antiguos:

La razón de la política recortada, R’, está restringida de tal manera que:

Dada la ventaja, At, como se muestra en la sección anterior, y la proporción de políticas, como se muestra arriba, la pérdida sustituta se calcula como:

El código a continuación muestra cómo implementar el mecanismo de recorte y la pérdida sustituta. 

def calculate_surrogate_loss(
        actions_log_probability_old,
        actions_log_probability_new,
        epsilon,
        advantages):
    advantages = advantages.detach()
    policy_ratio = (
            actions_log_probability_new - actions_log_probability_old
            ).exp()
    surrogate_loss_1 = policy_ratio * advantages
    surrogate_loss_2 = torch.clamp(
            policy_ratio, min=1.0-epsilon, max=1.0+epsilon
            ) * advantages
    surrogate_loss = torch.min(surrogate_loss_1, surrogate_loss_2)
    return surrogate_loss

3. Entrenando al Agente

Ahora, vamos a entrenar al agente.

Calculando la pérdida de política y valor 

Ahora estamos listos para calcular las pérdidas de política y valor: 

• La pérdida de política es la suma de la pérdida de sustituto y el bono de entropía.
• La pérdida de valor se basa en la diferencia entre el valor predicho por el crítico y los retornos (recompensa acumulativa) generados por la política. El cálculo de la pérdida de valor utiliza la función de pérdida Smooth L1. Esto ayuda a suavizar la función de pérdida y la hace menos sensible a los valores atípicos.

Ambas pérdidas, como se calcularon anteriormente, son tensores. El descenso de gradiente se basa en valores escalares. Para obtener un solo valor escalar que represente la pérdida, utiliza la función .sum() para sumar los elementos del tensor. La función a continuación muestra cómo hacer esto:

def calculate_losses(
        surrogate_loss, entropy, entropy_coefficient, returns, value_pred):
    entropy_bonus = entropy_coefficient * entropy
    policy_loss = -(surrogate_loss + entropy_bonus).sum()
    value_loss = f.smooth_l1_loss(returns, value_pred).sum()
    return policy_loss, value_loss

Definiendo el bucle de entrenamiento

Antes de comenzar el proceso de entrenamiento, crea un conjunto de buffers como arreglos vacíos. El algoritmo de entrenamiento utilizará estos buffers para almacenar información sobre las acciones del agente, los estados del entorno y las recompensas en cada paso de tiempo. La función a continuación inicializa estos buffers:

def init_training():
    states = []
    actions = []
    actions_log_probability = []
    values = []
    rewards = []
    done = False
    episode_reward = 0
    return states, actions, actions_log_probability, values, rewards, done, episode_reward

Cada iteración de entrenamiento ejecuta al agente con los parámetros de política para esa iteración. El agente interactúa con el entorno en pasos de tiempo en un bucle hasta que alcanza una condición terminal.

Después de cada paso de tiempo, la acción, recompensa y valor del agente se agregan a los buffers respectivos. Cuando el episodio termina, la función devuelve el conjunto actualizado de buffers, que resumen los resultados del episodio.

Antes de ejecutar el bucle de entrenamiento:

• Establece el modelo en modo de entrenamiento usando agent.train().
• Restablece el entorno a un estado aleatorio usando env.reset(). Este es el estado inicial para esta iteración de entrenamiento.

Los siguientes pasos explican qué sucede en cada paso de tiempo en el bucle de entrenamiento:

• Pasa el estado al agente.
• El agente devuelve:
• La acción predicha dada el estado, basada en la política (actor). Pasa este tensor de acción predicha a través de la función softmax para obtener el conjunto de probabilidades de acción.
• El valor predicho del estado, basado en el crítico.
• El agente selecciona la acción a tomar:
• Usa las probabilidades de acción para estimar la distribución de probabilidad.
• Selecciona aleatoriamente una acción tomando una muestra de esta distribución. La función dist.sample() hace esto.
• Utiliza la función env.step() para pasar esta acción al entorno y simular la respuesta del entorno para este paso de tiempo. Basado en la acción del agente, el entorno genera:
• El nuevo estado
• La recompensa
• El valor booleano de retorno hecho (esto indica si el entorno ha alcanzado un estado terminal)
• Agregar a los búferes respectivos los valores de la acción del agente, las recompensas, los valores predichos y el nuevo estado.

El episodio de entrenamiento termina cuando la función env.step() devuelve true para el valor booleano de retorno de hecho.

Después de que el episodio haya terminado, utiliza los valores acumulados de cada paso de tiempo para calcular los retornos acumulativos de este episodio sumando las recompensas de cada paso de tiempo. Utilizamos la función calculate_returns() descrita anteriormente para hacer esto. Los inputs de esta función son el factor de descuento y el buffer que contiene las recompensas de cada paso de tiempo. Utilizamos estos retornos y los valores acumulados de cada paso de tiempo para calcular las ventajas utilizando la función calculate_advantages().

La siguiente función en Python muestra cómo implementar estos pasos:

def forward_pass(env, agent, optimizer, discount_factor):
    states, actions, actions_log_probability, values, rewards, done, episode_reward = init_training()
    state = env.reset()
    agent.train()
    while not done:
        state = torch.FloatTensor(state).unsqueeze(0)
        states.append(state)
        action_pred, value_pred = agent(state)
        action_prob = f.softmax(action_pred, dim=-1)
        dist = distributions.Categorical(action_prob)
        action = dist.sample()
        log_prob_action = dist.log_prob(action)
        state, reward, done, _ = env.step(action.item())
        actions.append(action)
        actions_log_probability.append(log_prob_action)
        values.append(value_pred)
        rewards.append(reward)
        episode_reward += reward
    states = torch.cat(states)
    actions = torch.cat(actions)
    actions_log_probability = torch.cat(actions_log_probability)
    values = torch.cat(values).squeeze(-1)
    returns = calculate_returns(rewards, discount_factor)
    advantages = calculate_advantages(returns, values)
    return episode_reward, states, actions, actions_log_probability, advantages, returns

Actualización de los parámetros del modelo

Cada iteración de entrenamiento ejecuta el modelo a través de un episodio completo que consiste en muchos pasos de tiempo (hasta que alcanza una condición terminal). En cada paso de tiempo, almacenamos los parámetros de la política, la acción del agente, los retornos y las ventajas. Después de cada iteración, actualizamos el modelo basándonos en el rendimiento de la política a lo largo de todos los pasos de tiempo en esa iteración.

El número máximo de pasos de tiempo en el entorno CartPole es de 500. En entornos más complejos, hay más pasos de tiempo, incluso millones. En tales casos, el conjunto de datos de resultados del entrenamiento debe dividirse en lotes. El número de pasos de tiempo en cada lote se llama tamaño del lote de optimización.

Por lo tanto, los pasos para actualizar los parámetros del modelo son:

• Dividir el conjunto de datos de resultados del entrenamiento en lotes.
• Para cada lote:
• Obtener la acción del agente y el valor predicho para cada estado.
• Utilizar estas acciones predichas para estimar la nueva distribución de probabilidad de acciones.
• Usa esta distribución para calcular la entropía.
• Usa esta distribución para obtener la probabilidad logarítmica de las acciones en el conjunto de resultados de entrenamiento. Este es el nuevo conjunto de probabilidades logarítmicas de las acciones en el conjunto de resultados de entrenamiento. El antiguo conjunto de probabilidades logarítmicas de estas mismas acciones fue calculado en el bucle de entrenamiento explicado en la sección anterior.
• Calcula la pérdida sustituta utilizando las distribuciones de probabilidad antiguas y nuevas de las acciones.
• Calcula la pérdida de política y la pérdida de valor utilizando la pérdida sustituta, la entropía y las ventajas.
• Ejecuta .backward() por separado en las pérdidas de política y valor. Esto actualiza los gradientes en las funciones de pérdida.
• Ejecute .step() en el optimizador para actualizar los parámetros de la política. En este caso, utilizamos el optimizador Adam para equilibrar la velocidad y la robustez.
• Acumule las pérdidas de política y valor.
• Repita el pase hacia atrás (las operaciones anteriores) en cada lote varias veces, dependiendo del valor del parámetro PPO_STEPS. Repetir el pase hacia atrás en cada lote es computacionalmente eficiente porque aumenta efectivamente el tamaño del conjunto de datos de entrenamiento sin tener que ejecutar pases hacia adelante adicionales. El número de pasos de entorno en cada alternancia entre muestreo y optimización se llama tamaño del lote de iteración.
• Devuelva la pérdida promedio de política y la pérdida de valor.

El código a continuación implementa estos pasos:

def update_policy(
        agent,
        states,
        actions,
        actions_log_probability_old,
        advantages,
        returns,
        optimizer,
        ppo_steps,
        epsilon,
        entropy_coefficient):
    BATCH_SIZE = 128
    total_policy_loss = 0
    total_value_loss = 0
    actions_log_probability_old = actions_log_probability_old.detach()
    actions = actions.detach()
    training_results_dataset = TensorDataset(
            states,
            actions,
            actions_log_probability_old,
            advantages,
            returns)
    batch_dataset = DataLoader(
            training_results_dataset,
            batch_size=BATCH_SIZE,
            shuffle=False)
    for _ in range(ppo_steps):
        for batch_idx, (states, actions, actions_log_probability_old, advantages, returns) in enumerate(batch_dataset):
            # obtener nueva probabilidad logarítmica de acciones para todos los estados de entrada
            action_pred, value_pred = agent(states)
            value_pred = value_pred.squeeze(-1)
            action_prob = f.softmax(action_pred, dim=-1)
            probability_distribution_new = distributions.Categorical(
                    action_prob)
            entropy = probability_distribution_new.entropy()
            # estimar nuevas probabilidades logarítmicas usando las acciones antiguas
            actions_log_probability_new = probability_distribution_new.log_prob(actions)
            surrogate_loss = calculate_surrogate_loss(
                    actions_log_probability_old,
                    actions_log_probability_new,
                    epsilon,
                    advantages)
            policy_loss, value_loss = calculate_losses(
                    surrogate_loss,
                    entropy,
                    entropy_coefficient,
                    returns,
                    value_pred)
            optimizer.zero_grad()
            policy_loss.backward()
            value_loss.backward()
            optimizer.step()
            total_policy_loss += policy_loss.item()
            total_value_loss += value_loss.item()
    return total_policy_loss / ppo_steps, total_value_loss / ppo_steps

4. Ejecutando el Agente PPO 

Finalmente, ejecutemos el agente PPO.

Evaluando el rendimiento 

Para evaluar el rendimiento del agente, crea un entorno nuevo y calcula las recompensas acumulativas de ejecutar el agente en este nuevo entorno. Necesitas configurar el agente en modo de evaluación usando la función .eval(). Los pasos son los mismos que para el bucle de entrenamiento. El fragmento de código a continuación implementa la función de evaluación: 

def evaluate(env, agent):
    agent.eval()
    rewards = []
    done = False
    episode_reward = 0
    state = env.reset()
    while not done:
        state = torch.FloatTensor(state).unsqueeze(0)
        with torch.no_grad():
            action_pred, _ = agent(state)
            action_prob = f.softmax(action_pred, dim=-1)
        action = torch.argmax(action_prob, dim=-1)
        state, reward, done, _ = env.step(action.item())
        episode_reward += reward
    return episode_reward

Visualizando los resultados del entrenamiento

Utilizaremos la biblioteca Matplotlib para visualizar el progreso del proceso de entrenamiento. La función a continuación muestra cómo trazar las recompensas tanto de los bucles de entrenamiento como de los de prueba:

def plot_train_rewards(train_rewards, reward_threshold):
    plt.figure(figsize=(12, 8))
    plt.plot(train_rewards, label='Training Reward')
    plt.xlabel('Episode', fontsize=20)
    plt.ylabel('Training Reward', fontsize=20)
    plt.hlines(reward_threshold, 0, len(train_rewards), color='y')
    plt.legend(loc='lower right')
    plt.grid()
    plt.show()

def plot_test_rewards(test_rewards, reward_threshold):
    plt.figure(figsize=(12, 8))
    plt.plot(test_rewards, label='Testing Reward')
    plt.xlabel('Episode', fontsize=20)
    plt.ylabel('Testing Reward', fontsize=20)
    plt.hlines(reward_threshold, 0, len(test_rewards), color='y')
    plt.legend(loc='lower right')
    plt.grid()
    plt.show()

En los gráficos de ejemplo a continuación, mostramos las recompensas de entrenamiento y prueba, obtenidas aplicando la política en los entornos de entrenamiento y prueba respectivamente. Tenga en cuenta que la forma de estos gráficos aparecerá diferente cada vez que ejecute el código. Esto se debe a la aleatoriedad inherente al proceso de entrenamiento.

Recompensas de entrenamiento (obtenidas aplicando la política en el entorno de entrenamiento). Imagen por Autor.

Recompensas de prueba (obtenidas aplicando la política en el entorno de prueba). Imagen por Autor.

En los gráficos de salida mostrados arriba, observe el progreso del proceso de entrenamiento:

• La recompensa comienza desde valores bajos. A medida que avanza el entrenamiento, las recompensas aumentan.
• Las recompensas fluctúan aleatoriamente mientras aumentan. Esto se debe a que el agente explora el espacio de políticas.
• El entrenamiento termina, y las recompensas de prueba se han estabilizado alrededor del umbral (475) durante muchas iteraciones.
• Las recompensas están limitadas a 500. Estas son restricciones impuestas por el entorno (Gym CartPole v1).

De manera similar, puedes trazar los valores de pérdida y política a lo largo de las iteraciones:

def plot_losses(policy_losses, value_losses):
    plt.figure(figsize=(12, 8))
    plt.plot(value_losses, label='Value Losses')
    plt.plot(policy_losses, label='Policy Losses')
    plt.xlabel('Episode', fontsize=20)
    plt.ylabel('Loss', fontsize=20)
    plt.legend(loc='lower right')
    plt.grid()
    plt.show()

El gráfico de ejemplo a continuación muestra las pérdidas rastreadas a lo largo de los episodios de entrenamiento:

Pérdidas de valor y política a través del proceso de entrenamiento. Imagen por el Autor

Observa la gráfica y nota:

• Las pérdidas parecen estar distribuidas aleatoriamente y no siguen ningún patrón.
• Esto es típico del entrenamiento de RL, donde el objetivo no es minimizar la pérdida, sino maximizar las recompensas.

Ejecutar el algoritmo PPO

Ahora tienes todos los componentes para entrenar al agente usando PPO. Para ponerlo todo junto, necesitas:

• Declarar hiperparámetros como factor de descuento, tamaño de lote, tasa de aprendizaje, etc.
• Instanciar buffers como arreglos nulos para almacenar las recompensas y pérdidas de cada iteración.
• Crear una instancia de agente utilizando la función create_agent().
• Ejecutar iterativamente pases hacia adelante y hacia atrás utilizando las funciones forward_pass() y update_policy().
• Probar el rendimiento de la política utilizando la función evaluate().
• Agregar la política, las pérdidas de valor y las recompensas de las funciones de entrenamiento y evaluación a los buffers respectivos.
• Calcula el promedio de las recompensas y pérdidas de los últimos pasos. El ejemplo a continuación promedia las recompensas y pérdidas de los últimos 40 pasos de tiempo.
• Imprime los resultados de la evaluación cada ciertos pasos. El ejemplo a continuación imprime cada 10 pasos.
• Termina el proceso cuando el promedio de la recompensa cruza cierto umbral.

El código a continuación muestra cómo declarar una función que hace esto en Python:

def run_ppo():
    MAX_EPISODES = 500
    DISCOUNT_FACTOR = 0.99
    REWARD_THRESHOLD = 475
    PRINT_INTERVAL = 10
    PPO_STEPS = 8
    N_TRIALS = 100
    EPSILON = 0.2
    ENTROPY_COEFFICIENT = 0.01
    HIDDEN_DIMENSIONS = 64
    DROPOUT = 0.2
    LEARNING_RATE = 0.001
    train_rewards = []
    test_rewards = []
    policy_losses = []
    value_losses = []
    agent = create_agent(HIDDEN_DIMENSIONS, DROPOUT)
    optimizer = optim.Adam(agent.parameters(), lr=LEARNING_RATE)
    for episode in range(1, MAX_EPISODES+1):
        train_reward, states, actions, actions_log_probability, advantages, returns = forward_pass(
                env_train,
                agent,
                optimizer,
                DISCOUNT_FACTOR)
        policy_loss, value_loss = update_policy(
                agent,
                states,
                actions,
                actions_log_probability,
                advantages,
                returns,
                optimizer,
                PPO_STEPS,
                EPSILON,
                ENTROPY_COEFFICIENT)
        test_reward = evaluate(env_test, agent)
        policy_losses.append(policy_loss)
        value_losses.append(value_loss)
        train_rewards.append(train_reward)
        test_rewards.append(test_reward)
        mean_train_rewards = np.mean(train_rewards[-N_TRIALS:])
        mean_test_rewards = np.mean(test_rewards[-N_TRIALS:])
        mean_abs_policy_loss = np.mean(np.abs(policy_losses[-N_TRIALS:]))
        mean_abs_value_loss = np.mean(np.abs(value_losses[-N_TRIALS:]))
        if episode % PRINT_INTERVAL == 0:
            print(f'Episode: {episode:3} | \
                  Mean Train Rewards: {mean_train_rewards:3.1f} \
                  | Mean Test Rewards: {mean_test_rewards:3.1f} \
                  | Mean Abs Policy Loss: {mean_abs_policy_loss:2.2f} \
                  | Mean Abs Value Loss: {mean_abs_value_loss:2.2f}')
        if mean_test_rewards >= REWARD_THRESHOLD:
            print(f'Reached reward threshold in {episode} episodes')
            break
    plot_train_rewards(train_rewards, REWARD_THRESHOLD)
    plot_test_rewards(test_rewards, REWARD_THRESHOLD)
    plot_losses(policy_losses, value_losses)

Ejecuta el programa:

run_ppo()

La salida debería parecerse al ejemplo a continuación:

Episode:  10 | Mean Train Rewards: 22.3 | Mean Test Rewards: 30.4 | Mean Abs Policy Loss: 0.37 | Mean Abs Value Loss: 0.39 
Episode:  20 | Mean Train Rewards: 38.6 | Mean Test Rewards: 69.8 | Mean Abs Policy Loss: 0.46 | Mean Abs Value Loss: 0.37 
.
.
.
Episode: 100 | Mean Train Rewards: 289.5 | Mean Test Rewards: 427.3 | Mean Abs Policy Loss: 1.73 | Mean Abs Value Loss: 0.21 
Episode: 110 | Mean Train Rewards: 357.7 | Mean Test Rewards: 461.4 | Mean Abs Policy Loss: 1.86 | Mean Abs Value Loss: 0.22 
Reached reward threshold in 116 episodes

Puedes ver y ejecutar el programa de trabajo en este cuaderno DataLab!

5. Ajuste y Optimización de Hiperparámetros

En el aprendizaje automático, los hiperparámetros controlan el proceso de entrenamiento. A continuación, explico algunos de los hiperparámetros importantes utilizados en PPO:

• Tasa de aprendizaje: La tasa de aprendizaje decide cuánto pueden variar los parámetros de la política en cada iteración. En el descenso de gradiente estocástico, la cantidad por la cual se actualizan los parámetros de la política en cada iteración está decidida por el producto de la tasa de aprendizaje y el gradiente.
• Parámetro de recorte: Esto también se conoce como épsilon, ε. Decide el grado en que se recorta la proporción de políticas. Se permite que la proporción de las políticas nuevas y antiguas varíe en el rango [1-ε, 1+ε]. Cuando está fuera de este rango, se recorta artificialmente para que se encuentre dentro del rango.
• Tamaño del lote: Esto se refiere al número de pasos a considerar para cada actualización de gradiente. En PPO, el tamaño del lote es el número de pasos de tiempo necesarios para aplicar la política y calcular la pérdida sustituta para actualizar los parámetros de la política. En este artículo, utilizamos un tamaño de lote de 64. 
• Pasos de iteración: Este es el número de veces que cada lote se reutiliza para ejecutar el pase hacia atrás. El código en este artículo se refiere a esto como PPO_STEPS. En entornos complejos, correr el pase hacia adelante muchas veces es costoso computacionalmente. Una alternativa más eficiente es volver a ejecutar cada lote unas cuantas veces. Normalmente se recomienda usar un valor entre 5 y 10. 
• Factor de descuento: Esto también se conoce como gamma, γ. Expresa en qué medida las recompensas inmediatas son más valiosas que las recompensas futuras. Esto es similar al concepto de tasas de interés en el cálculo del valor temporal del dinero. Cuando está más cerca de 0, significa que las recompensas futuras son menos valiosas y el agente debe priorizar las recompensas inmediatas. Cuando está más cerca de 1, significa que las recompensas futuras son importantes.
• Coeficiente de entropía: El coeficiente de entropía decide el bono de entropía, que se calcula como el producto del coeficiente de entropía y la entropía de la distribución. El papel del bono de entropía es introducir más aleatoriedad en la política. Esto anima al agente a explorar el espacio de la política. Sin embargo, el entrenamiento no converge a una política óptima cuando esta aleatoriedad es demasiado alta.
• Criterios de éxito para el entrenamiento: Necesitas establecer los criterios para decidir cuándo el entrenamiento es exitoso. Una forma común de hacer esto es poner una condición de que las recompensas promedio en los últimos N intentos (episodios) estén por encima de cierto umbral. En el código de ejemplo anterior, esto se expresa con la variable N_TRIALS. Cuando se establece en un valor más alto, el entrenamiento lleva más tiempo porque la política debe lograr la recompensa objetivo en más episodios. También resulta en una política más robusta aunque computacionalmente más costosa. Ten en cuenta que PPO es una política estocástica, y habrá episodios en los que el agente no cruce el umbral. Por lo tanto, si el valor de N_TRIALS es demasiado alto, tu entrenamiento puede no finalizar.

Estrategias para optimizar el rendimiento de PPO

Optimizar el rendimiento de los algoritmos de entrenamiento PPO implica prueba y error y experimentación con diferentes valores de hiperparámetros. Sin embargo, hay algunas pautas generales:

• Factor de descuento: Cuando las recompensas a largo plazo son importantes, como en el entorno de CartPole, donde el poste necesita mantenerse estable con el tiempo, comience con un valor de gamma moderado, como 0.99.
• Bonificación de entropía: En entornos complejos, el agente debe explorar el espacio de acciones para encontrar la política óptima. La bonificación de entropía promueve la exploración. La bonificación de entropía se agrega a la pérdida sustituta. Verifique la magnitud de la pérdida sustituta y la entropía de la distribución antes de decidir el coeficiente de entropía. En este artículo, utilizamos un coeficiente de entropía de 0.01.
• Parámetro de recorte: El parámetro de recorte decide qué tan diferente puede ser la política actualizada de la política actual. Un valor grande del parámetro de recorte fomenta una mejor exploración del entorno, pero corre el riesgo de desestabilizar el entrenamiento. Se necesita un parámetro de recorte que permita una exploración gradual mientras evita actualizaciones que desestabilicen. En este artículo, utilizamos un parámetro de recorte de 0.2.
• Tasa de aprendizaje: Cuando la tasa de aprendizaje es demasiado alta, la política se actualiza en pasos grandes en cada iteración y el proceso de entrenamiento podría volverse inestable. Cuando es demasiado baja, el entrenamiento lleva demasiado tiempo. Este tutorial utilizó una tasa de aprendizaje de 0.001, que funciona bien para el entorno. En muchos casos, se recomienda utilizar una tasa de aprendizaje de 1e-5.

Desafíos y Mejores Prácticas en PPO

Después de explicar los conceptos y detalles de implementación de PPO, vamos a discutir los desafíos y las mejores prácticas.

Desafíos comunes en el entrenamiento de PPO

Aunque PPO se utiliza ampliamente, es necesario estar al tanto de los desafíos potenciales para resolver problemas del mundo real utilizando esta técnica con éxito. Algunos de estos desafíos son:

• Convergencia lenta: En entornos complejos, PPO puede ser ineficiente en cuanto a las muestras y necesita muchas interacciones con el entorno para converger en la política óptima. Esto hace que sea lento y costoso de entrenar.
• Sensibilidad a los hiperparámetros: PPO se basa en explorar eficientemente el espacio de políticas. La estabilidad del proceso de entrenamiento y la velocidad de convergencia son sensibles a los valores de los hiperparámetros. Los valores óptimos de estos hiperparámetros a menudo solo se pueden determinar mediante prueba y error. 
• Sobreajuste:  Los entornos de RL suelen inicializarse con parámetros aleatorios. El entrenamiento de PPO se basa en encontrar la política óptima basada en el entorno del agente. A veces, el proceso de entrenamiento converge a un conjunto de parámetros óptimos para un entorno específico pero no para cualquier entorno aleatorizado. Esto suele abordarse teniendo muchas iteraciones, cada una con un entorno de entrenamiento aleatorizado de manera diferente. 
• Ambientes dinámicos:Los ambientes de RL simples, como el ambiente CartPole, son estáticos: las reglas son las mismas a lo largo del tiempo. Muchos otros ambientes, como un robot que aprende a caminar sobre una superficie inestable en movimiento, son dinámicos: las reglas del ambiente cambian con el tiempo. Para desempeñarse bien en tales ambientes, PPO a menudo necesita ajustes adicionales.
• Exploración vs explotación:El mecanismo de recorte de PPO asegura que las actualizaciones de la política estén dentro de una región confiable. Sin embargo, también impide que el agente explore el espacio de acciones. Esto puede llevar a la convergencia a óptimos locales, especialmente en ambientes complejos. Por otro lado, permitir que el agente explore demasiado puede evitar que converja en cualquier política óptima.

Mejores prácticas para entrenar modelos PPO

Para obtener buenos resultados utilizando PPO, recomiendo algunas buenas prácticas, como:

• Normalizar las características de entrada: La normalización de los valores de retornos y ventajas reduce la variabilidad en los datos y conduce a actualizaciones de gradientes estables. Normalizar los datos lleva todos los valores a un rango numérico consistente. Ayuda a reducir el efecto de valores atípicos y extremos, que de otro modo podrían distorsionar las actualizaciones de gradientes y ralentizar la convergencia.
• Usar tamaños de lote adecuadamente grandes: Los lotes pequeños permiten actualizaciones y entrenamientos más rápidos, pero pueden llevar a la convergencia a óptimos locales e inestabilidad en el proceso de entrenamiento. Los tamaños de lote más grandes permiten al agente aprender políticas robustas, lo que conduce a un proceso de entrenamiento estable. Sin embargo, los tamaños de lote que son demasiado grandes también son subóptimos. Además de aumentar los costos computacionales, hacen que las actualizaciones de políticas sean menos receptivas a la función de valor porque las actualizaciones de gradiente se basan en promedios estimados sobre lotes grandes. Además, puede llevar a un sobreajuste de las actualizaciones a ese lote específico. 
• Pasos de iteración: Por lo general, es recomendable reutilizar cada lote durante 5-10 iteraciones. Esto hace que el proceso de entrenamiento sea más eficiente. Reutilizar el mismo lote demasiadas veces lleva al sobreajuste. El código se refiere a este hiperparámetro como PPO_STEPS. 
• Realizar evaluaciones regulares: Para detectar el sobreajuste, es esencial monitorear periódicamente la efectividad de la política. Si la política resulta ser ineficaz en ciertos escenarios, puede ser necesario realizar más entrenamiento o ajustes finos.
• Ajustar los hiperparámetros: Como se explicó anteriormente, el entrenamiento de PPO es sensible a los valores de los hiperparámetros. Experimenta con diversos valores de hiperparámetros para determinar el conjunto correcto de valores para tu problema específico.
• Red compartida de la red principal: Como se ilustra en este artículo, el uso de una red principal compartida evita desequilibrios entre las redes de actor y crítico. Compartir una red principal entre el actor y el crítico ayuda con la extracción de características compartidas y una comprensión común del entorno. Esto hace que el proceso de aprendizaje sea más eficiente y estable. También ayuda a reducir el espacio computacional y la complejidad temporal del algoritmo. 
• Número y tamaño de las capas ocultas: Aumente el número de capas ocultas y dimensiones para entornos más complejos. Problemas más simples como CartPole pueden resolverse con una sola capa oculta. La capa oculta utilizada en este artículo tiene 64 dimensiones. Hacer que la red sea mucho más grande de lo necesario es un desperdicio computacional y puede hacer que sea inestable. 
• Detención temprana: Detener el entrenamiento cuando se cumplen las métricas de evaluación ayuda a prevenir el sobreentrenamiento y evita el desperdicio de recursos. Una métrica de evaluación común es cuando el agente supera las recompensas de umbral en los últimos N eventos.

Conclusión

En este artículo, discutimos PPO como una forma de resolver problemas de RL. Luego detallamos los pasos para implementar PPO usando PyTorch. Por último, presentamos algunos consejos de rendimiento y mejores prácticas para PPO.

La mejor manera de aprender es implementar el código tú mismo. También puedes modificar el código para que funcione con otros entornos de control clásico en Gym. Para aprender cómo implementar agentes de RL usando Python y Gymnasium de OpenAI, ¡sigue el curso Reinforcement Learning with Gymnasium in Python!