Decolando para o sucesso: Lunar Lander com Python e Deep Q-Learning (DQN)

Prepare-se para uma aventura emocionante! Neste post, vamos dominar um desafio clássico do aprendizado por reforço: o Lunar Lander do OpenAI Gymnasium, utilizando a poderosa técnica de Deep Q-Learning (DQN) com Python e TensorFlow. Se você sempre quis ensinar um foguete virtual a pousar suavemente na Lua, este é o seu guia definitivo!

O que é o lunar lander e por que você deveria se importar?

Imagine que você é o controlador de voo de uma missão lunar, responsável por guiar um módulo lunar em segurança até a superfície. O Lunar Lander é um ambiente de simulação que replica esse cenário, onde seu objetivo é pousar o módulo entre as bandeirinhas sem causar um estrondo digno de Hollywood.

Por que se importar? Porque o Lunar Lander é um problema complexo que exige que um agente (nosso foguete) aprenda a tomar decisões em tempo real, baseando-se em uma série de observações (posição, velocidade, ângulo, etc.) e aplicando ações (ligar ou desligar os propulsores). Resolver esse problema demonstra a capacidade de criar agentes inteligentes que podem se adaptar e performar bem em ambientes desafiadores. E quem sabe, talvez você possa usar essas habilidades para construir um robô aspirador de pó que realmente funcione!

Deep Q-Learning (DQN): O cérebro por trás do foguete

DQN é um algoritmo de aprendizado por reforço que combina a técnica de Q-Learning com redes neurais profundas. Em termos simples, ele permite que um agente aprenda a tomar decisões ótimas em um ambiente, aprendendo a função Q, que estima a “qualidade” de uma ação em um determinado estado.

Pense na função Q como um conselheiro espacial que sussurra no ouvido do nosso foguete: “Se você fizer isso agora, terá uma boa chance de pousar suavemente (ou virar poeira estelar)”.

Como o DQN funciona na prática

1. Observação: O agente observa o estado atual do ambiente (posição, velocidade, etc.).
2. Ação: Baseado em sua função Q (implementada por uma rede neural), o agente escolhe a melhor ação a ser tomada. Às vezes, ele pode explorar ações aleatórias para descobrir novas estratégias (mais sobre isso em breve).
3. Recompensa: O agente recebe uma recompensa (positiva ou negativa) com base no resultado de sua ação. Um pouso suave rende uma recompensa alta, enquanto uma colisão explosiva resulta em uma recompensa negativa.
4. Aprendizado: O agente usa a recompensa para ajustar sua função Q, aprendendo quais ações são mais propensas a levar a um resultado positivo.
5. Repetição: O processo se repete até que o agente se torne um mestre em pousos lunares.

O código do sucesso: passo a passo

Agora, vamos mergulhar no código que torna tudo isso possível. Vou explicar cada parte em detalhes, para que você entenda exatamente o que está acontecendo nos bastidores.

1. Preparando o terreno (importando as bibliotecas)

Comece instalando as dependências necessárias:

pip install numpy tensorflow "gymnasium[all]"

Depois, importe as bibliotecas:

import random
from collections import deque

import gymnasium as gym
import numpy as np
import tensorflow as tf
from tensorflow.keras.layers import Dense
from tensorflow.keras.models import Sequential

• random: Para ações aleatórias durante a exploração.
• collections.deque: Uma fila otimizada para armazenar as experiências do agente.
• gymnasium as gym: O ambiente de simulação do Lunar Lander.
• numpy: Para manipulação eficiente de arrays.
• tensorflow e tensorflow.keras: Para construir e treinar a rede neural.

2. Definindo os hiperparâmetros: As chaves do nosso sucesso

# Define hiperparâmetros
learning_rate = 0.001
discount_factor = 0.99
epsilon_start = 1.0
epsilon_end = 0.01
epsilon_decay_steps = 10000
replay_buffer_size = 100000
batch_size = 64
sync_target_interval = 1000
weights_file = "lunarlander_dqn_weights.weights.h5"
load_existing_weights = True

Os hiperparâmetros são como os botões de controle da nossa máquina de aprendizado. Ajustá-los corretamente é crucial para o sucesso. Vamos entender cada um deles:

• learning_rate: A taxa de aprendizado da rede neural. Um valor muito alto pode levar a instabilidade, enquanto um valor muito baixo pode tornar o aprendizado muito lento. Imagine que é a velocidade com que o foguete assimila cada aprendizado.
• discount_factor: Um fator que determina a importância de recompensas futuras em relação às recompensas imediatas. Um valor próximo de 1 significa que o agente valoriza o longo prazo.
• epsilon_start e epsilon_end: Controlam a taxa de exploração do agente. No início, epsilon é alto, o que significa que o agente explora mais ações aleatórias para descobrir o ambiente. À medida que o treinamento avança, epsilon diminui, e o agente passa a explorar menos. Imagine o foguete no início, tentando entender o que acontece ao apertar cada botão do controle.
• epsilon_decay_steps: Define quantos passos leva para epsilon diminuir de epsilon_start para epsilon_end.
• replay_buffer_size: O tamanho do buffer de replay, onde as experiências do agente são armazenadas para treinamento.
• batch_size: O número de experiências que são amostradas do buffer de replay para cada atualização da rede neural.
• sync_target_interval: A frequência com que os pesos da rede target são atualizados com os pesos da rede q_network.
• weights_file: O nome do arquivo onde os pesos da rede neural serão salvos.
• load_existing_weights: Se deve carregar os pesos do arquivo salvo.

3. Criando o ambiente lunar lander

# Cria o ambiente LunarLander-v3
env = gym.make('LunarLander-v3')
state_space_size = env.observation_space.shape[-1]
action_space_size = env.action_space.n

Esta parte é simples: criamos o ambiente Lunar Lander usando a biblioteca gymnasium e obtemos o tamanho do espaço de estados (o número de variáveis que descrevem o estado do ambiente) e o tamanho do espaço de ações (o número de ações que o agente pode realizar).

4. Construindo a rede neural (Q-Network)

# Define a arquitetura da rede neural (Q-Network)
def create_q_network(state_space_size, action_space_size):
    model = Sequential([
        Dense(128, activation='relu', input_shape=(state_space_size,)),
        Dense(128, activation='relu'),
        Dense(action_space_size, activation='linear')
    ])
    model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=learning_rate), loss='mse')
    return model

# Cria as redes Q (online e target)
q_network = create_q_network(state_space_size, action_space_size)
target_network = create_q_network(state_space_size, action_space_size)
target_network.set_weights(q_network.get_weights())

Aqui, definimos a arquitetura da nossa rede neural, que será responsável por aproximar a função Q. Usamos um modelo sequencial com duas camadas densas com função de ativação ReLU, seguidas por uma camada de saída linear.

Por que duas redes? Utilizamos duas redes para treinar nosso agente: a q_network e a target_network. A q_network é atualizada a cada passo de treinamento, enquanto a target_network é atualizada apenas a cada sync_target_interval passos, recebendo os pesos da q_network. Isso ajuda a estabilizar o treinamento, pois impede que as atualizações da função Q oscilem muito rapidamente.

5. Implementando o agente DQN

# Implementa o agente DQN
class DQNAgent:
    def __init__(self, q_network, target_network, state_space_size, action_space_size, learning_rate, discount_factor, epsilon_start, epsilon_end, epsilon_decay_steps, replay_buffer_size, batch_size, sync_target_interval):
        self.state_space_size = state_space_size
        self.action_space_size = action_space_size
        self.learning_rate = learning_rate
        self.discount_factor = discount_factor
        self.epsilon = epsilon_start
        self.epsilon_start = epsilon_start
        self.epsilon_end = epsilon_end
        self.epsilon_decay_steps = epsilon_decay_steps
        self.replay_buffer = deque(maxlen=replay_buffer_size)
        self.batch_size = batch_size
        self.train_step = 0
        self.sync_target_interval = sync_target_interval
        self.q_optimizer = tf.keras.optimizers.Adam(learning_rate=learning_rate)
        self.q_network = q_network
        self.target_network = target_network

    def choose_action(self, state):
        if random.random() < self.epsilon:
            return env.action_space.sample()
        else:
            q_values = self._predict_q_values(np.expand_dims(state, axis=0))[0]
            return np.argmax(q_values)

    def store_experience(self, state, action, reward, next_state, done):
        self.replay_buffer.append((state, action, reward, next_state, done))

    def sample_experience(self):
        if len(self.replay_buffer) < self.batch_size:
            return None
        batch = random.sample(self.replay_buffer, self.batch_size)
        states, actions, rewards, next_states, dones = zip(*batch)
        return np.array(states), np.array(actions), np.array(rewards).reshape(-1, 1), np.array(next_states), np.array(dones).reshape(-1, 1)

    def update_q_network(self):
        if len(self.replay_buffer) < self.batch_size:
            return

        states, actions, rewards, next_states, dones = self.sample_experience()

        target_q_values = self._predict_target_q_values(next_states)
        max_next_q_values = tf.reduce_max(target_q_values, axis=1, keepdims=True)
        targets = rewards + self.discount_factor * max_next_q_values * (1 - dones)

        self._train_step(states, actions, targets)

        self.epsilon = self.epsilon_end + (self.epsilon_start - self.epsilon_end) * np.exp(
            -self.train_step / self.epsilon_decay_steps
        )
        self.train_step += 1

        if self.train_step % self.sync_target_interval == 0:
            self.target_network.set_weights(self.q_network.get_weights())

    @tf.function
    def _predict_q_values(self, state):
        return self.q_network(state)

    @tf.function
    def _train_step(self, states, actions, targets):
        with tf.GradientTape() as tape:
            q_values = self.q_network(states)
            action_masks = tf.one_hot(actions, depth=self.action_space_size)
            predicted_q_values = tf.reduce_sum(action_masks * q_values, axis=1, keepdims=True)
            loss = tf.keras.losses.mse(targets, predicted_q_values)
        gradients = tape.gradient(loss, self.q_network.trainable_variables)
        self.q_optimizer.apply_gradients(zip(gradients, self.q_network.trainable_variables))
        return loss

    @tf.function
    def _predict_target_q_values(self, next_states):
        return self.target_network(next_states)

Esta classe encapsula toda a lógica do agente DQN:

• __init__: Inicializa o agente com os hiperparâmetros e as redes neurais.
• choose_action: Escolhe uma ação com base na política epsilon-greedy. Com uma probabilidade epsilon, o agente escolhe uma ação aleatória (exploração). Caso contrário, ele escolhe a ação que maximiza a função Q (explotação).
• store_experience: Armazena a experiência do agente no buffer de replay.
• sample_experience: Amostra um lote de experiências aleatórias do buffer de replay.
• update_q_network: Atualiza a rede neural Q com base nas experiências amostradas do buffer de replay.

6. Treinamento do agente: A jornada do herói

if load_existing_weights:
    epsilon_start = 0.1

# Inicializa o agente
agent = DQNAgent(q_network, target_network, state_space_size, action_space_size, learning_rate, discount_factor, epsilon_start, epsilon_end, epsilon_decay_steps, replay_buffer_size, batch_size, sync_target_interval)

# Carrega pesos se solicitado
if load_existing_weights:
    try:
        q_network.load_weights(weights_file)
        target_network.load_weights(weights_file)
        print(f"Pesos carregados do arquivo: {weights_file}, Épsilon definido para: {agent.epsilon}")
    except FileNotFoundError:
        print("Arquivo de pesos não encontrado. Iniciando treinamento do zero.")

# Loop de treinamento
num_episodes = 5000  # Número de episódios
target_score = 200  # Recompensa média para considerar resolvido
consecutive_successes = 0
episode_rewards = deque(maxlen=100)  # Para calcular a média móvel

for episode in range(num_episodes):
    state, info = env.reset()
    total_reward = 0
    done = False

    while not done:
        action = agent.choose_action(state)
        next_state, reward, done, truncated, info = env.step(action)
        agent.store_experience(state, action, reward, next_state, done)
        agent.update_q_network()
        state = next_state
        total_reward += reward

    episode_rewards.append(total_reward)
    mean_reward = np.mean(episode_rewards)

    print(f"Episódio: {episode + 1}, Recompensa Total: {total_reward:.2f}, Recompensa Média (últimos 100): {mean_reward:.2f}, Épsilon: {agent.epsilon:.2f}")

    if mean_reward >= target_score and episode > 100:  # Espera um pouco antes de considerar resolvido
        consecutive_successes += 1
        if consecutive_successes >= 10:
            print(f"\nProblema resolvido em {episode + 1} episódios com recompensa média de {mean_reward:.2f}!")
            break
    else:
        consecutive_successes = 0

# Salva os pesos após o treinamento
q_network.save_weights(weights_file)
print(f"Pesos da rede Q salvos em: {weights_file}")

env.close()

Finalmente, chegamos ao loop de treinamento. Aqui, o agente interage com o ambiente Lunar Lander por um número definido de episódios. Em cada episódio, o agente observa o estado, escolhe uma ação, recebe uma recompensa e atualiza sua rede neural Q.

O treinamento continua até que o agente atinja uma recompensa média satisfatória, indicando que ele aprendeu a pousar o módulo lunar com sucesso.

Execute o código e acompanhe o treinamento. Inicialmente o módulo lunar vai ser muito errático, mas com o passar das épocas, o algoritmo vai entendendo o objetivo e pousando corretamente com mais frequência.

Se você não quiser passar pela etapa de treinamento, no repositório do Github tem o código fonte completo com o arquivo de pesos da rede pré treinada, para você conseguir ver resultados já na primeira execução.

Dicas e truques para um pouso perfeito

• Ajuste os Hiperparâmetros: Os hiperparâmetros certos podem fazer toda a diferença. Experimente diferentes valores para a taxa de aprendizado, o fator de desconto e a taxa de exploração para encontrar a combinação ideal.
• Explore o Espaço de Ações: Incentive o agente a explorar diferentes ações, especialmente no início do treinamento. Isso pode ajudá-lo a descobrir estratégias inesperadas.
• Monitore o Progresso: Acompanhe a recompensa média e a taxa de sucesso do agente para garantir que o treinamento esteja progredindo conforme o esperado.
• Tenha Paciência: O aprendizado por reforço pode levar tempo. Não desanime se o agente não aprender a pousar imediatamente. Continue treinando e ajustando os hiperparâmetros até que ele alcance o sucesso.

Sua missão começa agora!

E aí, preparado para decolar? Agora que você tem o conhecimento e o código, é hora de colocar a mão na massa e treinar seu próprio agente Lunar Lander. Compartilhe seus resultados, dúvidas e ideias nos comentários abaixo. E não se esqueça de compartilhar este post com seus amigos que também são apaixonados por inteligência artificial e aprendizado por reforço. Juntos, podemos conquistar a Lua (e outros desafios ainda maiores)!