Pular para o conteúdo principal

Primeiros passos com o TorchRL para aprendizado por reforço profundo

Um guia amigável para iniciantes sobre o TorchRL para aprendizado por reforço profundo - aprenda a criar agentes de RL com o PyTorch por meio de exemplos práticos.
Actualizado 30 de jan. de 2025  · 30 min de leitura

O aprendizado por reforço (RL) é usado para resolver muitos problemas complexos, desde o treinamento de carros autônomos até o treinamento de grandes modelos de linguagem (LLMs) para dar respostas semelhantes às humanas. 

Os agentes de RL são treinados para responder com base no feedback humano, usando o Reinforcement Learning from Human Feedback (RLHF). Embora estruturas baseadas em Python, como Keras e Tensorflow, sejam usadas em aplicativos empresariais de aprendizagem profunda, a maioria dos novos projetos agora se baseia no PyTorch e no PyTorch Lightning.

O TorchRL é uma biblioteca de código aberto para criar soluções de RL usando o PyTorch. Neste tutorial, mostrarei a você como configurar o TorchRL, entender seus componentes subjacentes e usá-lo para criar um agente de RL simples. Também discutiremos o uso do TorchRL para implementar versões pré-construídas de algoritmos de RL, como o Proximal Policy Optimization (PPO). Por fim, abordaremos os princípios básicos de registro e monitoramento de agentes de RL.

O que é o TorchRL e por que usá-lo?

Os algoritmos de RL geralmente são complexos. Você precisa criar o(s) agente(s), calcular os retornos e as perdas, definir as passagens para frente e para trás e avaliar o desempenho do agente. 

O TorchRL agrupa muitas funções de RL usadas com frequência em módulos que você pode acessar diretamente. Isso facilita a implementação e a experimentação de vários algoritmos para resolver problemas práticos. Isso também simplifica a criação de novos algoritmos, pois os pesquisadores têm acesso ao restante do ecossistema de RL sem precisar criá-lo por conta própria. 

O TorchRL vem com muitos módulos pré-construídos que tornam o desenvolvimento de RL mais eficiente. Por exemplo: 

  • Ambientes: O TorchRL fornece uma API padronizada para importar e usar ambientes de RL de várias fontes,incluindo Gymnasium, Jumanji, RoboHive e maisre. Em muitos casos, é necessário personalizar a saída do ambiente para atender a necessidades específicas de treinamento. O TorchRL inclui módulos para várias transformações de ambiente usando uma única chamada de função com os parâmetros desejados. 
  • Coletores de dados e buffers de repetição: Muitos algoritmos de treinamento de RL envolvem a coleta de dados sobre as interações do agente com o ambiente: a ação do agente, a recompensa recebida e o próximo estado em que ele se encontra. O TorchRL inclui pacotes com as estruturas de dados para coletar e coletar amostras dessas informações. 
  • Funções objetivas: A metodologia principal de um algoritmo de RL é implementada em sua função objetiva. O TorchRL empacota algoritmos comuns de RL como DQN (Deep Q Networks), A2C (Advantage Actor-Critic), PPO e muitos outros em módulos pré-construídos que você pode invocar diretamente e usar para treinar o agente. 

Dada a funcionalidade acima, o TorchRL provou ser útil para agilizar e simplificar a criação de soluções de RL para vários casos de uso, como:

  • Robótica: O robô precisa ser treinado para navegar com sucesso em um ambiente complexo. Você é penalizado por tropeçar ou não conseguir navegar. Usando métodos de treinamento de RL, ele aprende as ações corretas a serem executadas para navegar com sucesso em um ambiente como uma superfície natural irregular. 
  • IA do jogo: Os jogadores baseados em computador em jogos de vídeo e console precisam criar e improvisar seus movimentos em resposta às ações do jogador humano. A RL, com recompensas e penalidades, é usada para treinar esses jogadores a escolher os movimentos certos para jogar contra o humano. 
  • Sistemas autônomos: Os carros autônomos precisam navegar de forma independente em um ambiente complexo (como o tráfego rodoviário) e concluir a meta (chegar ao destino) sem acidentes. A RL é usada para treinar sistemas autônomos para imitar o comportamento de um motorista humano em diversas condições do mundo real. 

Desenvolver aplicativos de IA

Aprenda a criar aplicativos de IA usando a API OpenAI.
Comece a treinar gratuitamente

Configurando o TorchRL

Nesta seção, mostrarei a você como instalar e começar a usar o TorchRL. 

Pré-requisitos

Você precisa de alguns pacotes de software de dependência antes de instalar e usar o TorchRL:

  • PyTorch: O TorchRL é baseado no PyTorch, portanto, você precisa dele como pré-requisito. 
  • Ginásio: Você necessita do pacote Gymnasium para importar ambientes RL. A partir de janeiro de 2025, a versão mais recente do Gymnasium não é compatível com o TorchRL, conforme explicado nesta página de discussões do Git. Portanto, instale a versão mais antiga 0.29.1.
  • PyGame: É um pacote Python para videogames.  O TorchRL precisa dele para simular ambientes de RL semelhantes a jogos, como o CartPole. 
  • TensorDict: O uso de uma estrutura de dados semelhante a um dicionário para armazenar as entradas e saídas das redes neurais torna conveniente trabalhar com tensores no TorchRL. O TensorDict fornece um contêiner de tensor que armazena tensores como pares de valores-chave. 

Instale os pacotes de pré-requisitos:

!pip install torch tensordict  gymnasium==0.29.1 pygame

Instalando o TorchRL

Instale o TorchRL usando pip. Recomendo fortemente a instalação desses pacotes em um ambiente Conda se você estiver trabalhando em um computador pessoal ou servidor. 

!pip install torchrl

Verificação da instalação

Após a instalação, verifique se o TorchRL foi instalado com êxito. Tente importar torchrl em um shell Python (ou notebook). Use o método check_env_specs() para verificar se um ambiente padrão (como o CartPole) corresponde às especificações de torchrl:

import torchrl 
from torchrl.envs import GymEnv 
from torchrl.envs.utils import check_env_specs

check_env_specs(GymEnv("CartPole-v1"))

O resultado deve indicar que você criou o ambiente com sucesso e que ele funciona com o TorchRL. 

[torchrl][INFO] check_env_specs succeeded!

Componentes principais do TorchRL

Antes de criar seu primeiro agente de RL, vamos examinar os principais componentes do TorchRL.

Ambientes

O TorchRL fornece uma API uniforme para interagir com diferentes ambientes. Isso é feito por meio do agrupamento de funcionalidades específicas de cada ambiente em um conjunto padrão de classes e funções de agrupamento. Você passa os parâmetros apropriados para o wrapper e ele mapeia internamente seu comando para a chamada de função correspondente ao ambiente específico. Em particular:

  • O TorchRL converte estados do ambiente (observações), ações e recompensas em objetos tensores do PyTorch, que podem ser usados diretamente pelos módulos que implementam algoritmos de RL.
  • Ele possibilita a aplicação de etapas de pré e pós-processamento, por exemplo, para normalizar e dimensionar tensores de entrada ou para apresentar tensores de saída em um formato específico.

Por exemplo, para criar um ambiente a partir do Gymnasium, use o módulo GymEnv:

env = GymEnv("CartPole-v1")

Transformações

É comum ampliar o ambiente de estoque com recursos adicionais e transformações que, de outra forma, você teria que implementar por conta própria. Por exemplo, você pode adicionar um módulo de contador de passos ao ambiente em vez de codificar o contador por conta própria. O contador de passos mantém o controle do número de passos em cada episódio. O módulo transformedEnv ajuda você com isso:

from torchrl.envs import GymEnv, StepCounter, TransformedEnv
env = TransformedEnv(GymEnv("CartPole-v1"), StepCounter())

Da mesma forma, a normalização de tensores antes de operá-los é uma etapa comum de pré-processamento. Uma transformação de normalização, usando o módulo ObservationNorm , normaliza os tensores. A documentação da função TransformedEnv discute vários tipos de transformações. 

Você também pode combinar várias transformações usando o parâmetro compose():

base_env = GymEnv('CartPole-v1', device=device) 

env = TransformedEnv( 
    base_env, 
    Compose(
        ObservationNorm(in_keys=["observation"]), 
        StepCounter()
    )
)

Agentes e políticas

Na RL, o agente decide suas ações usando a política e com base no estado observado do ambiente. O objetivo do agente é maximizar as recompensas cumulativas do ambiente. Ele recebe recompensas quando escolhe a ação correta (como parar em um sinal vermelho). 

Por exemplo, em um carro autônomo, o agente dirige o carro. Suas decisões (como dirigir o carro em uma determinada direção) baseiam-se em suas observações do estado do ambiente (como tráfego, posição do carro etc.) e de sua política (como evitar pedestres e outros obstáculos, parar em sinais vermelhos etc.). 

A política mais simples é escolher uma ação aleatoriamente. O ator escolhe aleatoriamente uma das ações do espaço de ações possíveis. Às vezes, uma política aleatória é usada para gerar um conjunto de dados inicial de interações antes de começar a treinar o modelo. 

Use o módulo RandomPolicy para criar uma política aleatória. Essa política aleatória aceita um parâmetro action_spec que especifica o espaço de ação. No exemplo abaixo, o espaço de ação consiste em números (contínuos) entre -1 e +1. Portanto, a política aleatória escolhe uma ação representada por um número aleatório entre -1 e +1. 

import torchrl 
import torch
from tensordict import TensorDict
from torchrl.data.tensor_specs import Bounded

action_spec = Bounded(-torch.ones(1), torch.ones(1))
actor = torchrl.envs.utils.RandomPolicy(action_spec=action_spec) 
td = actor(TensorDict({}, batch_size=[])) 
print(td.get("action"))

O resultado deve ser um tensor, como mostrado abaixo:

tensor([0.9258])

Verifique se a política é de fato aleatória executando o ator novamente:

td = actor(TensorDict({}, batch_size=[])) 
print(td.get("action"))

Ele deve produzir um tensor aleatório diferente.

Se você precisar de uma introdução aos conceitos de RL, confira a trilha de habilidades do Reinforcement Learning in Python no DataCamp!

Criando seu primeiro agente de RL com o TorchRL

Nesta seção, mostrarei a você como implementar um agente simples de Aprendizado por reforço usando o TorchRL. 

Antes de começar, importe os pacotes de software de pré-requisito no Python:

  • time para medir o tempo necessário para treinar o agente.
  • GymEnv Você pode usar os sites StepCounter e TransformedEnv para trabalhar em ambientes de ginásio.
  • MLP para criar uma rede neural MLP (perceptron de várias camadas) simples. 
  • EGreedyModule para equilibrar a exploração do ambiente e a exploração da política mais conhecida.
  • QValueModule e DQNLoss para implementar o algoritmo Deep Q-Learning.
  • SoftUpdate para atualizar a rede neural.
  • SyncDataCollector para coletar dados das interações do agente. 
  • ReplayBuffer para armazenar os dados das interações do agente.
  • Adam para retropropagação.
  • matplotlib para exibir visualmente o progresso do treinamento. 
  • torchrl_logger para registrar a sessão de treinamento.
import time
import matplotlib.pyplot as plt
from torchrl.envs import GymEnv, StepCounter, TransformedEnv
from tensordict.nn import TensorDictModule as TensorDict, TensorDictSequential as Seq
from torchrl.modules import EGreedyModule, MLP, QValueModule
from torchrl.objectives import DQNLoss, SoftUpdate
from torchrl.collectors import SyncDataCollector
from torchrl.data import LazyTensorStorage, ReplayBuffer
from torch.optim import Adam
from torchrl._utils import logger as torchrl_logger

Etapa 1: Definir o ambiente

Neste exemplo, resolvemos o ambiente CartPole. Importe esse ambiente do Gymnasium junto com um contador de passos para acompanhar o número de passos de treinamento:

env = TransformedEnv(GymEnv("CartPole-v1"), StepCounter())

Semeie os ambientes Python e RL para replicar resultados semelhantes nas sessões de treinamento. 

torch.manual_seed(0)
env = TransformedEnv(GymEnv("CartPole-v1"), StepCounter())
env.set_seed(0)

Defina os parâmetros e hiperparâmetros para o treinamento: 

  • INIT_RAND_STEPS: O número de etapas para as quais o agente age aleatoriamente antes de usar a política. Essas etapas iniciais são usadas para coletar dados iniciais para começar a treinar a política.
  • FRAMES_PER_BATCH: O número de pontos de dados (um para cada interação ou etapa de tempo) em um lote de treinamento. 
  • OPTIM_STEPS: O número de etapas para as quais você deve acumular as perdas antes de executar uma passagem da passagem para trás. 
  • EPS_0: O valor inicial de epsilon, o coeficiente de exploração. 
  • BUFFER_LEN: O tamanho do buffer de repetição.
  • ALPHA: A taxa de aprendizado.
  • TARGET_UPDATE_EPS: O fator de decaimento para atualizar a rede de destino uusa o módulo de atualização suave.
  • REPLAY_BUFFER_SAMPLE: O tamanho da amostra aleatória a ser selecionada do buffer de repetição em cada iteração de treinamento. O buffer de reprodução armazena os resultados das interações do agente com o ambiente em cada intervalo de tempo. 
  • LOG_EVERY: O número de etapas após as quais você deve imprimir o progresso do treinamento.
  • MLP_SIZE: O tamanho da rede neural.
INIT_RAND_STEPS = 5000 
FRAMES_PER_BATCH = 100
OPTIM_STEPS = 10
EPS_0 = 0.5
BUFFER_LEN = 100_000
ALPHA = 0.05
TARGET_UPDATE_EPS = 0.95
REPLAY_BUFFER_SAMPLE = 128
LOG_EVERY = 1000
MLP_SIZE = 64

Etapa 2: Criar a política

Defina uma rede neural simples para implementar a política: 

  • Defina a MLP (rede neural). Com base em uma observação, o MLP gera o valor da ação, realizando efetivamente o trabalho da função Q. 
  • Defina um dicionário de tensores usando o módulo TensorDict no MLP acima. Esse dicionário mapeia as observações de estado do ambiente (chaves do dicionário) para os valores de probabilidade de ação (valores do dicionário) correspondentes a esse estado. 
  • Use o QValueModule para implantar a função Q. Dado um tensor que contém valores de ação, ele retorna a ação correspondente ao valor de ação mais alto. Ele implementa a estratégia gananciosa (exploração). 
  • Combine (usando omódulo TensorDictSequential) o dicionário de tensores do MLP e o site QValueModule para definir a política .
value_mlp = MLP(out_features=env.action_spec.shape[-1], num_cells=[MLP_SIZE, MLP_SIZE])
value_net = TensorDict(value_mlp, in_keys=["observation"], out_keys=["action_value"])
policy = Seq(value_net, QValueModule(spec=env.action_spec))
  • Defina a função de exploração usando o EGreedyModule - ele recebe como entrada o espaço de ação do ambiente, o comprimento do buffer de reprodução e o coeficiente de exploração, epsilon. Encadeie (usando o módulo TensorDictSequential ) essa função de exploração com a política definida acima para obter a política final:
exploration_module = EGreedyModule(
    env.action_spec, annealing_num_steps=BUFFER_LEN, eps_init=EPS_0
)
policy_explore = Seq(policy, exploration_module)

Etapa 3: Treinar o agente

A primeira etapa do treinamento do agente é coletar os dados das interações do agente com o ambiente. Use o SyncDataCollector para criar um coletor para executar a política e coletar os resultados das interações do agente: 

collector = SyncDataCollector(
    env,
    policy_explore,
    frames_per_batch=FRAMES_PER_BATCH,
    total_frames=-1,
    init_random_frames=INIT_RAND_STEPS,
)

Crie um buffer de repetição para armazenar os resultados das interações:

rb = ReplayBuffer(storage=LazyTensorStorage(BUFFER_LEN))

Você também precisa declarar módulos específicos de treinamento, como a função de perda (para usar o cálculo de perda baseado em DQN), o otimizador (o algoritmo tradicional de Adam) e as funções de atualização (para atualizar a rede neural). Todos eles são baseados em módulos TorchRL predefinidos:

loss = DQNLoss(value_network=policy, action_space=env.action_spec, delay_value=True)
optim = Adam(loss.parameters(), lr=ALPHA)
updater = SoftUpdate(loss, eps=TARGET_UPDATE_EPS)

Inicialize os contadores para manter o controle do número total de etapas e episódios, das etapas bem-sucedidas por episódio e do tempo de execução:

total_count = 0
total_episodes = 0
t0 = time.time()
success_steps = []

A função de treinamento consiste em 2 loops for:

  • O loop superior executa a política e anexa os resultados das interações do agente ao buffer de reprodução.
  • O loop interno divide as amostras de treinamento em lotes. Ele processa cada lote da seguinte forma: 
    • Escolha uma amostra aleatória do buffer de reprodução.
    • Calcule as perdas usando a função de perda.
    • Execute o otimizador e o atualizador.
    • Atualize os contadores.

Cada etapa envolve a chamada de módulos TorchRL pré-construídos sem que você precise codificar nada do zero. O código abaixo mostra como você pode implementar essas etapas:

for i, data in enumerate(collector):
    rb.extend(data)
    max_length = rb[:]["next", "step_count"].max()
    if len(rb) > INIT_RAND_STEPS:
        for _ in range(OPTIM_STEPS):
            sample = rb.sample(REPLAY_BUFFER_SAMPLE)
            loss_vals = loss(sample)
            loss_vals["loss"].backward()
            optim.step()

            optim.zero_grad()
            # Update exploration factor
            exploration_module.step(data.numel())

            # Update target params
            updater.step()
            total_count += data.numel()
            total_episodes += data["next", "done"].sum()
    success_steps.append(max_length)

Etapa 4: Avaliar o agente

O loop da seção anterior treina continuamente a política. Precisamos definir os critérios para avaliar o desempenho e decidir quando considerar o treinamento bem-sucedido. Também queremos que você tenha acesso ao progresso do treinamento. 

Imprimimos o progresso do treinamento em intervalos periódicos usando o registrador TorchRL:

if total_count > 0 and total_count % LOG_EVERY == 0:
    torchrl_logger.info(f"Successful steps in the last episode: {max_length}, rb length {len(rb)}, Number of episodes: {total_episodes}")

Usamos o número máximo de etapas que o agente alcançou no último episódio para determinar se o treinamento foi bem-sucedido. O ambiente do CartPole-v1 limita o número máximo de etapas e a recompensa total por episódio a 500. É comum considerar uma política bem-sucedida se ela atingir mais de 475 etapas:

if max_length > 475:
    print("TRAINING COMPLETE")
    break

Os dois trechos de código acima devem ser anexados ao final do loop de treinamento (mostrado anteriormente). O trecho a seguir mostra o loop de treinamento com o código para avaliar o agente: 

for i, data in enumerate(collector):
    # Write data in replay buffer
    rb.extend(data)
    max_length = rb[:]["next", "step_count"].max()
    if len(rb) > INIT_RAND_STEPS:
        for _ in range(OPTIM_STEPS):
            sample = rb.sample(REPLAY_BUFFER_SAMPLE)
            loss_vals = loss(sample)
            loss_vals["loss"].backward()
            optim.step()

            optim.zero_grad()
            # Update exploration factor
            exploration_module.step(data.numel())

            # Update target params
            updater.step()
            total_count += data.numel()
            total_episodes += data["next", "done"].sum()
    success_steps.append(max_length)

    if total_count > 0 and total_count % LOG_EVERY == 0:
        torchrl_logger.info(f"Successful steps in the last episode: {max_length}, rb length {len(rb)}, Number of episodes: {total_episodes}")

    if max_length > 475:
        print("TRAINING COMPLETE")
        break

Por fim, após o treinamento e a avaliação, imprima o tempo total de treinamento e trace o progresso do treinamento:

t1 = time.time()

torchrl_logger.info(
    f"solved after {total_count} steps, {total_episodes} episodes and in {t1-t0}s."
)

def plot_steps():
    plt.plot(success_steps)
    plt.title('Successful steps over training episodes')
    plt.xlabel('Training episodes')
    plt.ylabel('Steps')
    plt.show()

plot_steps()

Você pode encontrar e executar o código completo para implementar o DQN usando o TorchRL nesta pasta de trabalho do DataLab

Nesta seção, criamos e treinamos um agente de RL usando o algoritmo DQN simples. Na próxima seção, exploraremos como usar módulos TorchRL pré-construídos para implementar um algoritmo mais complexo, como o PPO. 

Explorando algoritmos pré-construídos no TorchRL

Conforme mencionado nas seções anteriores, o TorchRL vem com alguns algoritmos pré-construídos. Vamos dar uma olhada neles e em como funcionam.

Algoritmos suportados

O TorchRL inclui módulos pré-construídos para muitos algoritmos comuns de Deep Reinforcement Learning, como: 

  • Redes Q profundas (DQN)
  • Gradiente de política determinística profunda (DDPG)
  • Ator-crítico suave (SAC)
  • Aprendizado duplo Q aleatório (REDQ)
  • CrossQ
  • Aprendizado implícito de Q (IQL)
  • Aprendizado contínuo de Q (CQL)
  • Aprendizagem por imitação adversarial generativa (GAIL)
  • Transformador de decisão (DT)
  • DDPG com atraso duplo (TD3) 
  • Crítica da vantagem do ator (A2C)
  • Otimização da política proximal (PPO)
  • REINFORCE
  • e mais 

Isso torna eficiente a experimentação de diferentes tipos de algoritmos e o estudo do desempenho de cada um deles para resolver um determinado problema. 

Exemplo: PPO na TorchRL

O Proximal Policy Optimization usa uma função de objetivo substituto recortado para obter um processo de treinamento suave que equilibra a exploração e o aproveitamento. 

Nesta seção, mostrarei como usar o módulo PPO pré-construído no TorchRL. De forma semelhante à implementação anterior, importe os módulos de pré-requisito. Além dos módulos (como ReplayBuffer, SyncDataCollector, etc.) necessários para a implementação anterior, você precisa de alguns pacotes adicionais para o PPO:

  • ProbabilisticActor, para escolher uma ação de forma estocástica. A escolha da ação de forma estocástica (em vez de sempre seguir a política e maximizar a recompensa) facilita a exploração do ambiente e a descoberta de caminhos mais ideais durante o treinamento.
  • OneHotCategorical, para gerar uma codificação one-hot do tensor que denota as probabilidades de log das ações. 
  • ValueOperator, para criar um módulo TorchRL com base na rede neural que implementa a função de valor.
  • GAE, para implementar a estimativa de vantagem generalizada. O PPO usa uma função de vantagem como um substituto para a função de valor. O módulo de valor (acima) é uma entrada para o GAE. 
  • ClipPPOLoss, para criar a função objetiva recortada para implementar o PPO. 
import torch
from torch import nn

from torchrl.envs import Compose, ObservationNorm, DoubleToFloat, StepCounter, TransformedEnv
from torchrl.envs.libs.gym import GymEnv
from torchrl.envs.utils import check_env_specs
from torchrl.modules import ProbabilisticActor, OneHotCategorical, ValueOperator
from torchrl.collectors import SyncDataCollector
from torchrl.data.replay_buffers import ReplayBuffer
from torchrl.data.replay_buffers.storages import LazyTensorStorage
from torchrl.data.replay_buffers.samplers import SamplerWithoutReplacement
from torchrl.objectives.value import GAE
from torchrl.objectives import ClipPPOLoss
from tensordict.nn import TensorDictModule

Declare os parâmetros e hiperparâmetros para o treinamento:

  • FRAMES_PER_BATCH: O número de quadros (intervalos de tempo) em cada lote de dados resultantes das interações do agente com o ambiente. 
  • SUB_BATCH_SIZE: O número de intervalos de tempo em cada sub-lote. No treinamento de PPO, cada lote é dividido em sub-lotes. 
  • TOTAL_FRAMES: O número total de intervalos de tempo para os quais você deve executar o treinamento. 
  • GAMMA: O fator de desconto para descontar o valor da recompensa de etapas de tempo futuras.
  • CLIP_EPSILON: Isso define a região de confiança dentro da qual a política pode mudar em cada iteração de treinamento subsequente. Ela é implementada cortando a política atualizada de modo que a proporção entre as probabilidades novas e antigas fique dentro de um determinado limite. 
  • ALPHA: A taxa de aprendizado.
  • ENTROPY_EPS: Esse hiperparâmetro controla a proporção de exploração e aproveitamento. 
  • OPTIM_STEPS: O número de vezes que você deve executar o otimizador em cada sub-batelada de dados. 
  • LOG_EVERY: O número de etapas após as quais você deve avaliar a política e imprimir as recompensas. 
FRAMES_PER_BATCH = 1024
TOTAL_FRAMES = 1048576
GAMMA = 0.99
LAMBDA = 0.95
CLIP_EPSILON = 0.2
ALPHA = 1e-4
ENTROPY_EPS = 5e-4 
SUB_BATCH_SIZE = 64
OPTIM_STEPS = 8
LOG_EVERY = 16

Importe o ambiente CartPole básico do Gymnasium:

device="cpu"
base_env = GymEnv('CartPole-v1', device=device) 

Declare o ambiente TorchRL importando o ambiente básico e adicionando módulos para normalizar os tensores de observação e contar o número de etapas:

env = TransformedEnv( 
    base_env, 
    Compose(
        ObservationNorm(in_keys=["observation"]), 
        DoubleToFloat(), 
        StepCounter()
    )
)

Inicialize o ambiente e semeie os geradores de números aleatórios:

env.transform[0].init_stats(1024) 
torch.manual_seed(0)
env.set_seed(0)
check_env_specs(env) 

Declare o ator como uma rede neural com 1 camada oculta e 16 pesos. Ele recebe como entrada a observação do ambiente e gera a probabilidade de cada ação no espaço de ação. 

actor_net = nn.Sequential(
    nn.Linear(env.observation_spec["observation"].shape[-1], 32, device=device),
    nn.ReLU(),
    nn.Linear(32, 32, device=device),
    nn.ReLU(),
    nn.Linear(32, env.action_spec.shape[-1], device=device),
    nn.ReLU()
)

Crie um dicionário de tensores (com base na rede neural acima) com chaves e valores que mapeiam os estados para as probabilidades de ação. 

actor_module = TensorDictModule(actor_net, in_keys=["observation"], out_keys=["logits"])

Para treinar o PPO (e muitos outros algoritmos de RL), você nem sempre escolhe a ação que leva à maior recompensa na próxima etapa. Ter algum grau de aleatoriedade na escolha da ação permite que o agente explore o ambiente e descubra caminhos melhores que podem levar a retornos de longo prazo mais altos. Crie um ator probabilístico para fazer isso:

actor = ProbabilisticActor(
    module = actor_module,
    spec = env.action_spec,
    in_keys = ["logits"],
    distribution_class = OneHotCategorical, 
    return_log_prob = True
)

Crie a rede para implementar a função de valor. Essa rede tem 16 pesos. Ele recebe como entrada o estado do ambiente (observação) e gera o valor esperado desse estado. 

value_net = nn.Sequential(
    nn.Linear(env.observation_spec["observation"].shape[-1], 16, device=device),
    nn.ReLU(),
    nn.Linear(16, 1, device=device),
    nn.ReLU()
)

Crie um módulo TorchRL que envolva a função de valor (implementada acima) usando o ValueOperator(). Esse módulo faz interface com outros componentes do TorchRL. 

value_module = ValueOperator(
    module = value_net,
    in_keys = ["observation"]
)

Crie o buffer de reprodução para armazenar os resultados das interações do agente com o ambiente: 

replay_buffer = ReplayBuffer(
    storage = LazyTensorStorage(max_size=FRAMES_PER_BATCH),
    sampler = SamplerWithoutReplacement()
)

Crie um coletor de dados para executar a política no ambiente e coletar os resultados das interações do agente com o ambiente em cada intervalo de tempo (quadro):

collector = SyncDataCollector(
    env,
    actor,
    frames_per_batch = FRAMES_PER_BATCH,
    total_frames = TOTAL_FRAMES,
    split_trajs = True,
    reset_at_each_iter = True,
    device=device
)

Use o módulo GAE para implementar a função de vantagem para PPO. A função de vantagem é baseada na função de valor. 

advantage_module = GAE(
    gamma = GAMMA, 
    lmbda = GAMMA, 
    value_network = value_module,
    average_gae = True
) 

Use o módulo integrado ClipPPOLoss para implementar a função de perda de acordo com o algoritmo PPO:

loss_module = ClipPPOLoss(
    actor_network = actor,
    critic_network = value_module,
    clip_epsilon = CLIP_EPSILON,
    entropy_bonus = bool(ENTROPY_EPS),
    entropy_coef = ENTROPY_EPS
)

Declare o otimizador Adam: 

optim = torch.optim.Adam(loss_module.parameters(), lr=ALPHA)

Crie um agendador para reduzir gradualmente a taxa de aprendizado à medida que o treinamento avança: 

scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optim, TOTAL_FRAMES // FRAMES_PER_BATCH)

Crie o loop de treinamento usando os componentes criados acima. A função de treinamento consiste em três loops for:

  • O loop mais externo percorre os dados coletados das interações do agente com o ambiente. O parâmetro TOTAL_FRAMES decide quantos intervalos de tempo você deve executar o agente .
  • O loop intermediário faz o seguinte: 
    • Ele executa o loop mais interno por um número fixo de iterações (normalmente entre 5 e 10), conforme definido pelo hiperparâmetro OPTIM_STEPS .
    • Ele atualiza a taxa de aprendizado após cada iteração de treinamento. 
    • Ele também calcula o valor da função de vantagem para cada intervalo de tempo. 
    • Avalie a política periodicamente e imprima as recompensas. 
  • O loop mais interno executa o loop de treinamento:
    • Amostragem de um lote de dados de treinamento do buffer de reprodução.
    • Calcule a perda usando o módulo de perda. 
    • Execute o backpropagation na perda.
    • Use o otimizador para fazer a descida do gradiente. 

O código abaixo implementa o loop de treinamento:

for i, tensordict_data in enumerate(collector): 
    for _ in range(OPTIM_STEPS): 
        advantage_module(tensordict_data)
        replay_buffer.extend(tensordict_data.reshape(-1).cpu())
        for _ in range(FRAMES_PER_BATCH // SUB_BATCH_SIZE): 
            data = replay_buffer.sample(SUB_BATCH_SIZE)
            loss = loss_module(data.to(device))
            loss_value = loss["loss_objective"] + loss["loss_critic"] + loss["loss_entropy"]
            loss_value.backward()
            optim.step()
            optim.zero_grad()
    scheduler.step()

    if i % LOG_EVERY == 0:
        with torch.no_grad():
            rollout = env.rollout(FRAMES_PER_BATCH, actor)
            reward_eval = rollout["next","reward"].sum()
            #print(rollout["next","reward"].sum())
            print(reward_eval)
            rewards.append(reward_eval)
            del rollout

Essa pasta de trabalho do DataLab tem o código (como mostrado acima) para treinar um agente de RL usando o algoritmo PPO do TorchRL. Use-o como ponto de partida para ajustar os parâmetros e melhorar o desempenho do agente. 

Personalização de algoritmos

O TorchRL tem um design modular que possibilita a personalização das soluções RL. A estrutura cuida da interface entre os diferentes componentes. Para testar e comparar seu desempenho, você pode conectar e reproduzir ambientes, políticas diferentes, buffers de reprodução e algoritmos de RL. Devido a essa modularidade:

  • Você pode trocar diferentes componentes sem precisar reescrever grandes partes do programa. 
  • Os componentes individuais também podem ser modificados de forma independente para criar soluções personalizadas. 
  • Você pode fazer experiências com diferentes módulos e criar uma solução otimizada. 

Alguns exemplos de personalização no TorchRL são:

  • Os loops de treinamento podem ser personalizados com programadores de taxa de aprendizado, registradores e métricas personalizadas. 
  • Diferentes comprimentos de buffers de repetição podem ser usados para ajustar a duração do treinamento.
  • O ambiente pode ser personalizado para incluir funções de escala e normalização, contadores de etapas, etc. 

Treinamento em visualização e depuração de RL

As seções anteriores mostraram como usar os módulos TorchRL para treinar agentes de RL. Nesta seção, discutiremos maneiras de monitorar e visualizar o progresso do treinamento. 

Monitoramento do progresso do treinamento

É útil registrar várias métricas durante o processo de treinamento para monitorar seu progresso. Pacotes como o TensorBoard facilitam o registro direto dos resultados durante o processo de treinamento. Por exemplo, o SummaryWriter do TensorBoard permite que você o chame dentro do loop de treinamento para registrar as métricas. O pseudocódigo a seguir mostra como você pode fazer isso:

from torch.utils.tensorboard import SummaryWriter
writer = SummaryWriter(log_dir="training_logs")

for _ in training_data (): # training loop
    loss = … 
    loss.backprop() 
    writer.add_scalar(“Loss - “, loss.item(), step)
   … 

writer.close() 

Após concluir o treinamento e o registro, você pode traçar e visualizar os resultados. Por exemplo, use o comando tensorboard para visualizar o progresso do treinamento .

tensorboard --logdir=”training_logs”

Depuração com o TorchRL

A depuração é essencial para verificar a interação do agente com o ambiente. Como primeira etapa, você deseja verificar o espaço de ação e o espaço de observação. A documentação do ambiente deve incluir as especificações do ambiente. Por exemplo, a documentação do CartPole mostra que: 

  • O espaço de ação inclui dois valores discretos - 0 (para empurrar o carrinho para a esquerda) e 1 (para empurrar o carrinho para a direita) 
  • O espaço de observação tem 4 valores: 
    • Posição do carrinho, com um valor entre 
    • -4,8 e +4,8. 
    • Velocidade do carrinho, com qualquer valor de número real
    • Ângulo do polo, com um valor entre -24° e +24°.
    • Velocidade angular do polo, com qualquer valor de número real

Inspecione os espaços de observação e ação para verificar se eles correspondem aos valores esperados. Por exemplo:

print("Observation space: ", env.observation_spec)
print("Action space: ", env.action_spec)

Além disso, você pode extrair amostras aleatórias dos espaços de observação e ação e verificar seus valores: 

print("Sample Observation:", base_env.observation_spec.sample().get("observation"))
print("Sample Action:", base_env.action_spec.sample())

O resultado é semelhante ao exemplo abaixo:

Sample Observation: tensor([-4.5960e+00,  3.4028e+38,  2.2261e-02,  3.4028e+38])
Sample Action: tensor([1, 0])

Observe que os 4 valores de estado estão dentro dos intervalos especificados anteriormente. Repita os comandos acima algumas vezes e observe que os diferentes valores de saída (amostrados aleatoriamente) estão dentro dos intervalos especificados. 

Os dois comandos acima são baseados no ambiente básico, importado diretamente do Gymnasium. Na prática, usamos o ambiente transformado, que aplica transformações, como a normalização, nos valores originais do tensor. Assim, quando você usa o ambiente TorchRL transformado para extrair uma amostra de observações, seus valores podem não estar mais dentro do mesmo intervalo que no ambiente Gymnasium original. 

Por exemplo, extraia uma amostra do ambiente transformado:  

print("Sample Observation:", env.observation_spec.sample().get("observation"))

O resultado é semelhante: 

Sample Observation: tensor([-57.5869,      nan,   4.5276,      nan])

Observe que os valores da posição do carrinho e do ângulo do polo estão fora dos intervalos. Isso ocorre porque os valores do tensor foram normalizados. Além disso, a velocidade do carrinho e a velocidade angular do polo têm valores nan porque o agente ainda não começou a interagir com o ambiente .

Visualização do desempenho do agente

Além de traçar o progresso do treinamento, pode ser útil renderizar o ambiente e observar visualmente as interações do agente. Isso pode fornecer informações sobre o desempenho do agente. 

A maneira mais pragmática de visualizar o ambiente é renderizando um vídeo. Você precisa de alguns pacotes adicionais: 

  • torchvisionVocê pode trabalhar com arquivos multimídia
  • avVocê pode usar o pacote ffmpeg para fazer a interface com o PyTorch .
pip install torchvision
pip install av==12.0.0

Após concluir o treinamento, siga estas etapas preparatórias para renderizar o vídeo:

  • Declare o caminho (relativo) para armazenar o vídeo.
  • Inicialize o registrador para armazenar os resultados das interações do agente no formato CSV.
  • Inicialize o gravador de vídeo para usar os registros CSV para gerar um vídeo.
  • Transforme o ambiente do TorchRL para incluir o gravador de vídeo.
…
# training loop
# for _ in …
…
…

from torchrl.record import CSVLogger, VideoRecorder 

path = "./training_loop"
logger = CSVLogger(exp_name="dqn", log_dir=path, video_format="mp4")
video_recorder = VideoRecorder(logger, tag="video")
record_env = TransformedEnv(
    GymEnv("CartPole-v1", from_pixels=True, pixels_only=False), video_recorder
)

Execute a política (treinada) no ambiente e despeje as renderizações no arquivo de vídeo:

record_env.rollout(max_steps=1000, policy=policy)
video_recorder.dump()

Depois de executar o programa, você encontrará o vídeo no caminho do diretório (./training_loop no trecho acima) que você definiu anteriormente. Observe emque a pasta de trabalho do DataLab para implementar o DQN usando o TorchRL não inclui o código para criar o vídeo porque é difícil exportar arquivos de vídeo com um notebook on-line. Execute o programa (e adicione as etapas acima) em um computador ou servidor local para gravar vídeos.

Práticas recomendadas para usar o TorchRL

Por fim, vamos abordar algumas práticas recomendadas. Aqui estão minhas recomendações.

Comece com ambientes simples

Apesar da facilidade de desenvolvimento graças ao TorchRL, o treinamento de agentes de RL para que tenham um bom desempenho em ambientes complexos continua sendo um desafio. Portanto, antes de tentar resolver problemas difíceis, é essencial resolver ambientes simples como o CartPole. 

A modularidade do TorchRL permite a experimentação de diferentes algoritmos e parâmetros. Antes de usar os algoritmos alternativos e as opções de personalização na prática, é necessário que você tenha algumas percepções sobre seu funcionamento. A melhor maneira de fazer isso é em um ambiente simples, em que os efeitos das alterações individuais sejam fáceis de observar visualmente. 

Experimente os hiperparâmetros

O desempenho do treinamento dos modelos de RL é sensível a vários hiperparâmetros, como taxa de aprendizado, taxa de exploração, taxa de desconto e outros hiperparâmetros específicos do algoritmo, como a taxa de corte para PPO. Por exemplo, uma taxa de aprendizado muito alta torna o treinamento instável, enquanto uma taxa de aprendizado muito baixa leva muito tempo para convergir. 

Da mesma forma, uma taxa de exploração muito alta impede a convergência do treinamento, enquanto uma taxa de exploração muito baixa pode impedir que o agente descubra o caminho ideal. 

Não existe uma fórmula para determinar os valores corretos dos hiperparâmetros. Existem diretrizes e valores recomendados para ambientes padrão, mas eles podem não se aplicar a todos os problemas. Portanto, é necessário experimentar vários métodos, como pesquisa em grade ou pesquisa aleatória, para determinar os melhores valores para os hiperparâmetros. 

Também é possível usar bibliotecas automatizadas como Weights & Biases Sweeps ou Optuna para testaro desempenho do treinamento em várias combinações de hiperparâmetros .

Aproveitar os algoritmos pré-construídos

Como você viu nos exemplos anteriores, usar os módulos pré-construídos do TorchRL economiza um esforço considerável de desenvolvimento. A alternativa é criar toda a funcionalidade do zero, o que consumiria muito mais tempo e custo para desenvolver e testar. 

Como muitos desenvolvedores usam os módulos do TorchRL, eles também se beneficiam do fato de terem sido testados em vários cenários. Assim, você pode esperar que eles sejam mais livres de bugs do que um módulo personalizado. Portanto, para a maioria dos casos de uso padrão, é altamente recomendável que você use módulos pré-construídos. 

Conclusão

Neste artigo, abordamos os conceitos básicos do TorchRL, uma estrutura baseada no PyTorch para implementar algoritmos de RL. Também vimos exemplos práticos de uso do TorchRL para implementar soluções mais simples, como o Deep Q-Learning, e algoritmos mais complexos, como o PPO. 

Na próxima etapa, você pode usar esses programas como base para fazer experiências com outros ambientes e algoritmos. 

Se você deseja aprofundar seu conhecimento sobre os fundamentos e as aplicações práticas do RL, confira Reinforcement Learning with Gymnasium in Python para obter mais experiência prática com ambientes baseados no Gymnasium.

Conceitos de inteligência artificial (IA) em Python

Comece a usar a IA

Arun Nanda's photo
Author
Arun Nanda
LinkedIn

Arun é um ex-fundador de startup que gosta de criar coisas novas. Atualmente, ele está explorando os fundamentos técnicos e matemáticos da Inteligência Artificial. Ele adora compartilhar o que aprendeu, por isso escreve sobre isso.

Além do DataCamp, você pode ler as publicações dele no Medium, Airbyte e Vultr.

Temas

Saiba mais sobre o aprendizado por reforço com estes cursos!

curso

Deep Reinforcement Learning in Python

4 hr
1.5K
Learn and use powerful Deep Reinforcement Learning algorithms, including refinement and optimization techniques.
Ver DetalhesRight Arrow
Iniciar curso
Ver maisRight Arrow
Relacionado

tutorial

Uma introdução ao Q-Learning: Um tutorial para iniciantes

Saiba mais sobre o algoritmo mais popular de aprendizado por reforço sem modelo com um tutorial em Python.
Abid Ali Awan's photo

Abid Ali Awan

16 min

tutorial

Guia de torchchat do PyTorch: Configuração local com Python

Saiba como configurar o torchchat do PyTorch localmente com Python neste tutorial prático, que fornece orientação e exemplos passo a passo.
François Aubry's photo

François Aubry

tutorial

Como treinar um LLM com o PyTorch

Domine o processo de treinamento de grandes modelos de linguagem usando o PyTorch, desde a configuração inicial até a implementação final.
Zoumana Keita 's photo

Zoumana Keita

8 min

tutorial

Criando agentes LangChain para automatizar tarefas em Python

Um tutorial abrangente sobre a criação de agentes LangChain com várias ferramentas para automatizar tarefas em Python usando LLMs e modelos de bate-papo usando OpenAI.
Bex Tuychiev's photo

Bex Tuychiev

14 min

tutorial

Criando um transformador com o PyTorch

Saiba como criar um modelo Transformer usando o PyTorch, uma ferramenta avançada de aprendizado de máquina moderno.
Arjun Sarkar's photo

Arjun Sarkar

26 min

tutorial

Dominando a retropropagação: Um guia abrangente para redes neurais

Mergulhe nos fundamentos da retropropagação em redes neurais com um guia prático para treinar e avaliar um modelo para um cenário de uso de classificação de imagens.
Zoumana Keita 's photo

Zoumana Keita

14 min

Ver maisVer mais