Guia de como utilizar a biblioteca Stable Baselines para projetos de Aprendizado por Reforço.
O guia completo encontra-se neste notebook:
Entretanto, caso você não queira ver os exemplos em código, o texto do guia está presente a seguir.
A Stable Baselines é uma biblioteca de Aprendizado por Reforço que implementa diversos algoritmos de agentes de RL, além de várias funcionalidades úteis para o treinamento de um agente. Suas implementações são bem simples e intuitivas, mas sem deixarem de ser otimizadas e poderosas, buscando facilitar o desenvolvimento de projetos de Aprendizado por Reforço de alta qualidade.
Para instalar a biblioteca com o pip, basta rodar:
pip install stable-baselines3
OBS: É necessário instalar antes o PyTorch.
Com Stable Baselines, o processo de criar e treinar um agente é bem simples. Entretanto, caso você não saiba muito de Aprendizado por Reforço, é primeiro preciso passar por alguns conhecimentos básicos.
O Gym é uma biblioteca desenvolvida pela OpenAI que contém várias implementações prontas de ambientes de Aprendizado por Reforço. Ela é muito utilizada quando se quer testar um algoritmo de agente sem ter o trabalho de programar seu próprio ambiente.
Um Ambiente de Aprendizado por Reforço é um espaço que representa o nosso problema, é o objeto com o qual o nosso agente deve interagir para cumprir sua função. Isso significa que o agente toma ações nesse ambiente, e recebe recompensas dele com base na qualidade de sua tomada de decisões.
Todos os ambientes são dotados de um espaço de observações, que é a forma pela qual o agente recebe informações e deve se basear para a tomada de decisões, e um espaço de ações, que especifica as ações possíveis do agente. No xadrez, por exemplo, o espaço de observações seria o conjunto de todas as configurações diferentes do tabuleiro, e o espaço de ações seria o conjunto de todos os movimentos permitidos.
Agora que você já sabe o que é um ambiente, é preciso entender como nosso agente interage efetivamente com ele. Todos os ambientes do Gym possuem alguns métodos simples para facilitar a comunicação com eles:
| Método | Funcionalidade |
|---|---|
reset() |
Inicializa o ambiente e recebe a observação inicial |
step(action) |
Executa uma ação e recebe a observação e a recompensa |
render() |
Renderiza o ambiente |
close() |
Fecha o ambiente |
Assim, o código para interagir com o ambiente costuma seguir o seguinte modelo:
ambiente = gym.make("Nome do Ambiente") # Cria o ambiente
observação = ambiente.reset() # Inicializa o ambiente
acabou = False
while not acabou:
ambiente.render() # Renderiza o ambiente
observação, recompensa, acabou, info = ambiente.step(ação) # Executa uma ação
ambiente.close() # Fecha o ambientePara utilizar um dos ambientes do Gym, nós utilizamos a função gym.make(), passando o nome do ambiente desejado como parâmetro e guardando seu valor retornado em uma variável que chamaramos de env. A lista com todos os ambiente pode ser encontrada aqui.
env = gym.make("CartPole-v1")Para exemplificar, vamos pensar no ambiente CartPole-v1, um ambiente bem simples que modela um pêndulo invertido em cima de um carrinho buscando seu estado de equilíbrio.
Antes de treinar qualquer agente, primeiro é preciso entender melhor quais as características do nosso ambiente.
O Espaço de Observação do CartPole é definido por 4 informações:
| Informação | Min | Max | |
|---|---|---|---|
| 0 | Posição do Carrinho | -4.8 | 4.8 |
| 1 | Velocidade do Carrinho | -Inf | Inf |
| 2 | Ângulo da Barra | -24 deg | 24 deg |
| 3 | Velocidade na Extremidade da Barra | -Inf | Inf |
Dessa forma, a cada instante recebemos uma lista da observação com o seguinte formato:
[-3.3715708e+00 -2.6997593e+38 2.7833584e-01 2.0276438e+38]Já o Espaço de Ação é composto por duas ações únicas: mover o carrinho para a esquerda ou para a direita.
Quando queremos mover o carrinho para a esquerda, fazemos um env.step(0); quando queremos movê-lo para a direita, enviamos um env.step(1)
Depois de escolhermos nosso ambiente, já podemos pensar em qual algoritmo de agente queremos usar.
A biblioteca disponibiliza algoritmos de diversos tipos, como Policy Gradients, Actor Critics, DQN, etc. Nem todos eles suportam todos os tipos de ambientes, então é recomendável dar uma olhada na página oficial dos algoritmos.
Todos os algoritmos são inicializados de uma forma parecida, nós instanciamos eles com alguns parâmetros em comum: policy, que define a arquitetura da rede neural e env, que define o ambiente no qual o agente vai treinar. Assim, a inicialização segue o seguinte formato:
agente = ALGORITMO(policy, env)Como exemplo, vamos criar um PPO, um tipo de Actor-Critic:
from stable_baselines import PPO
model = PPO('MlpPolicy', env, seed=1, verbose=1)Todos os agentes também possuem alguns métodos em comuns bem importantes de se conhecer:
| Método | Funcionalidade |
|---|---|
learn() |
Treina o agente |
predict(obs) |
Escolhe uma ação com base na observação |
save(caminho) |
Salve o agente |
load(caminho) |
Carrega o agente |
Dessa forma, se quisermos escolher a próxima ação do nosso agente, nós rodamos:
agente.predict(observação)Bom, agora que já temos nosso agente e nosso ambiente, já podemos rodar nosso primeiro episódio!
Para isso, vamos criar uma função run_episode para simplificar o processo:
# A função recebe o ambiente e o agente como parâmetros
def run_episode(env, model, render=False):
# Primeiro, inicializamos o ambiente e guardamos a observação inicial em 'obs'
obs = env.reset()
# Guarda se o episódio terminou ou não
done = False
# Loop do episódio
while not done:
# Nosso modelo prediz a ação 'action' a ser tomada com base na nossa observação 'obs'
action, _states = model.predict(obs)
# Tomamos a ação 'action', e recebemos uma nova observação 'obs', uma recompensa 'reward'
# e se o episódio terminou 'done'
obs, reward, done, info = env.step(action)
# Renderiza o ambiente, caso desejado
if render:
env.render()
# Finaliza o episódio, caso tenha terminado
if done:
break
# Quando terminado, fechamos o ambiente
env.close()Para rodar o episódio, basta fazer:
run_episode(env, model, render=True)Se você tentou rodar um episódio, provavelmente o resultado não foi tão bom assim. Isso é porque precisamos treinar nosso agente para que ele saiba as melhores ações a se tomar.
Para melhor avaliar o desempenho do nosso agente, podemos utilizar a função evaluate_policy da biblioteca, que roda uma quantidade determinável de episódios e retorna a recompensa média obtida.
# Ambiente separado para avaliação
eval_env = gym.make('CartPole-v1')
# Avaliando o agente
mean_reward, std_reward = evaluate_policy(model, eval_env, n_eval_episodes=25, deterministic=True)O treinamento do agente acontece de maneira bem simples, basta rodar o método .learn() com a quantidade de instantes de tempo total_timesteps que desejamos treinar.
Ao longo do treinamento, nosso agente vai mostrando algumas informações importantes no ouput, como a duração média dos episódios mean_episode_length, a entropia entropy, o custo da função de custo loss, etc.
model.learn(total_timesteps=20020)Depois de treinarmos o nosso agente, podemos rodar mais um episódio para ver como ele melhorou.
run_episode(env, model, render=True)Esse é todo o conhecimento necessário para resolver mais ambientes simples. Entretanto, ainda existem várias outras funcionalidades muito interessantes da biblioteca que valem a pena aprender.
Para obter mais informações do treinamento, podemos utilizar o wrapper Monitor da biblioteca para monitorar o desempenho do nosso agente mesmo durante o treino.
Para isto, primeiro devemos criar uma pasta de logs:
import os
# Cria um diretório de logs
log_dir = "./logs/"
os.makedirs(log_dir, exist_ok=True)Em seguida, criamos o nosso ambiente e passamos ele para o nosso wrapper:
from stable_baselines3.common.monitor import Monitor
# Cria o ambiente
env = gym.make('CartPole-v1')
# Encapsula ele no wrapper Monitor
env = Monitor(env, log_dir)A partir daí, basta treinar o nosso agente como normal, utilizando o nosso novo ambiente encapsulado pelo wrapper:
model = PPO('MlpPolicy', env, seed=1, verbose=1).learn(total_timesteps=40000)Terminado o treinamento, podemos plotar os resultados utilizando o result_plotter:
from stable_baselines3.common import results_plotter
# Plota os resultados
results_plotter.plot_results([log_dir], 1e5, results_plotter.X_TIMESTEPS, "PPO CartPole")
Pronto! Agora temos uma visualização da performance do nosso agente durante o seu treinamento. Entretanto, podemos criar a nossa própria função de visualização para melhor analisar o modelo:
import matplotlib.pyplot as plt
def plot_results(log_folder, window=10, title='Curva de Aprendizado'):
"""
Plota os resultados.
:param log_folder: (str) diretório dos resultados a serem plotados
:param window: (int) tamanho da janela da média móvel
:param title: (str) título do plot
"""
# Obtém os resultados
results = results_plotter.load_results(log_folder)
x, y = results_plotter.ts2xy(results, 'timesteps')
# Calcula a média móvel do retorno
y_smoothed = results_plotter.rolling_window(y, window=window).mean(axis=1)
# Plota os resultados
fig = plt.figure(title, figsize=(10, 5))
plt.scatter(x, y, s=2)
plt.plot(x[window-1:], y_smoothed, color='darkblue', label="Média Móvel")
plt.xlabel('Timesteps')
plt.ylabel('Retorno Médio')
plt.title(title)
plt.legend()
plt.show()Agora conseguimos visualizar melhor o treinamento:
plot_results(log_dir)