DQN C51 / Rainbow

Ver no TensorFlow.org Executar no Google Colab Ver fonte no GitHubBaixar caderno

Introdução

Este exemplo mostra como treinar um agente DQN categórico (C51) no ambiente Cartpole usando a biblioteca TF-Agents.

Ambiente de cartpole

Certifique-se de dar uma olhada no tutorial DQN como um pré-requisito. Este tutorial presumirá familiaridade com o tutorial DQN; ele se concentrará principalmente nas diferenças entre DQN e C51.

Configurar

Se você ainda não instalou tf-agents, execute:

sudo apt-get update
sudo apt-get install -y xvfb ffmpeg
pip install 'imageio==2.4.0'
pip install pyvirtualdisplay
pip install tf-agents
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import base64
import imageio
import IPython
import matplotlib
import matplotlib.pyplot as plt
import PIL.Image
import pyvirtualdisplay

import tensorflow as tf

from tf_agents.agents.categorical_dqn import categorical_dqn_agent
from tf_agents.drivers import dynamic_step_driver
from tf_agents.environments import suite_gym
from tf_agents.environments import tf_py_environment
from tf_agents.eval import metric_utils
from tf_agents.metrics import tf_metrics
from tf_agents.networks import categorical_q_network
from tf_agents.policies import random_tf_policy
from tf_agents.replay_buffers import tf_uniform_replay_buffer
from tf_agents.trajectories import trajectory
from tf_agents.utils import common

tf.compat.v1.enable_v2_behavior()


# Set up a virtual display for rendering OpenAI gym environments.
display = pyvirtualdisplay.Display(visible=0, size=(1400, 900)).start()

Hiperparâmetros

env_name = "CartPole-v1" # @param {type:"string"}
num_iterations = 15000 # @param {type:"integer"}

initial_collect_steps = 1000  # @param {type:"integer"} 
collect_steps_per_iteration = 1  # @param {type:"integer"}
replay_buffer_capacity = 100000  # @param {type:"integer"}

fc_layer_params = (100,)

batch_size = 64  # @param {type:"integer"}
learning_rate = 1e-3  # @param {type:"number"}
gamma = 0.99
log_interval = 200  # @param {type:"integer"}

num_atoms = 51  # @param {type:"integer"}
min_q_value = -20  # @param {type:"integer"}
max_q_value = 20  # @param {type:"integer"}
n_step_update = 2  # @param {type:"integer"}

num_eval_episodes = 10  # @param {type:"integer"}
eval_interval = 1000  # @param {type:"integer"}

Meio Ambiente

Carregue o ambiente como antes, com um para treinamento e outro para avaliação. Aqui usamos CartPole-v1 (vs. CartPole-v0 no tutorial DQN), que tem uma recompensa máxima maior de 500 em vez de 200.

train_py_env = suite_gym.load(env_name)
eval_py_env = suite_gym.load(env_name)

train_env = tf_py_environment.TFPyEnvironment(train_py_env)
eval_env = tf_py_environment.TFPyEnvironment(eval_py_env)

Agente

C51 é um algoritmo de Q-learning baseado em DQN. Como o DQN, ele pode ser usado em qualquer ambiente com um espaço de ação discreto.

A principal diferença entre C51 e DQN é que em vez de simplesmente prever o valor Q para cada par de ação de estado, C51 prevê um modelo de histograma para a distribuição de probabilidade do valor Q:

Exemplo de distribuição C51

Ao aprender a distribuição ao invés de simplesmente o valor esperado, o algoritmo é capaz de ficar mais estável durante o treinamento, levando a um melhor desempenho final. Isso é particularmente verdadeiro em situações com distribuições de valores bimodais ou mesmo multimodais, onde uma única média não fornece uma imagem precisa.

Para treinar em distribuições de probabilidade em vez de em valores, C51 deve realizar alguns cálculos de distribuição complexos para calcular sua função de perda. Mas não se preocupe, tudo isso é cuidado para você no TF-Agents!

Para criar um Agente C51, primeiro precisamos criar uma CategoricalQNetwork . A API de CategoricalQNetwork é a mesma da QNetwork , exceto que há um argumento adicional num_atoms . Isso representa o número de pontos de suporte em nossas estimativas de distribuição de probabilidade. (A imagem acima inclui 10 pontos de suporte, cada um representado por uma barra azul vertical.) Como você pode ver pelo nome, o número padrão de átomos é 51.

categorical_q_net = categorical_q_network.CategoricalQNetwork(
    train_env.observation_spec(),
    train_env.action_spec(),
    num_atoms=num_atoms,
    fc_layer_params=fc_layer_params)

Também precisamos de um optimizer para treinar a rede que acabamos de criar e uma variável train_step_counter para controlar quantas vezes a rede foi atualizada.

Observe que uma outra diferença significativa do vanilla DqnAgent é que agora precisamos especificar min_q_value e max_q_value como argumentos. Eles especificam os valores mais extremos do suporte (em outras palavras, o mais extremo dos 51 átomos de cada lado). Certifique-se de escolhê-los apropriadamente para o seu ambiente específico. Aqui usamos -20 e 20.

optimizer = tf.compat.v1.train.AdamOptimizer(learning_rate=learning_rate)

train_step_counter = tf.compat.v2.Variable(0)

agent = categorical_dqn_agent.CategoricalDqnAgent(
    train_env.time_step_spec(),
    train_env.action_spec(),
    categorical_q_network=categorical_q_net,
    optimizer=optimizer,
    min_q_value=min_q_value,
    max_q_value=max_q_value,
    n_step_update=n_step_update,
    td_errors_loss_fn=common.element_wise_squared_loss,
    gamma=gamma,
    train_step_counter=train_step_counter)
agent.initialize()

Uma última coisa a notar é que também adicionamos um argumento para usar atualizações de n-passos com $ n $ = 2. No Q-learning de uma única etapa ($ n $ = 1), calculamos apenas o erro entre os valores-Q na etapa de tempo atual e na próxima etapa de tempo usando o retorno de etapa única (com base na equação de otimização de Bellman). O retorno de uma única etapa é definido como:

$ G_t = R_ {t + 1} + \ gama V (s_ {t + 1}) $

onde definimos $ V (s) = \ max_a {Q (s, a)} $.

As atualizações de N etapas envolvem expandir a função de retorno de etapa única padrão $ n $ vezes:

$ G_t ^ n = R_ {t + 1} + \ gamma R_ {t + 2} + \ gamma ^ 2 R_ {t + 3} + \ dots + \ gamma ^ n V (s_ {t + n}) $

As atualizações em N etapas permitem que o agente inicialize a partir de um futuro mais distante e, com o valor correto de $ n $, isso geralmente leva a um aprendizado mais rápido.

Embora as atualizações do C51 e de n etapas sejam frequentemente combinadas com a reprodução priorizada para formar o núcleo do agente Rainbow , não vimos nenhuma melhoria mensurável na implementação da reprodução priorizada. Além disso, descobrimos que, ao combinar nosso agente C51 apenas com atualizações de n etapas, nosso agente tem um desempenho tão bom quanto outros agentes Rainbow na amostra de ambientes Atari que testamos.

Métricas e Avaliação

A métrica mais comum usada para avaliar uma política é o retorno médio. O retorno é a soma das recompensas obtidas durante a execução de uma política em um ambiente para um episódio, e normalmente fazemos a média disso em alguns episódios. Podemos calcular a métrica de retorno médio da seguinte maneira.

def compute_avg_return(environment, policy, num_episodes=10):

  total_return = 0.0
  for _ in range(num_episodes):

    time_step = environment.reset()
    episode_return = 0.0

    while not time_step.is_last():
      action_step = policy.action(time_step)
      time_step = environment.step(action_step.action)
      episode_return += time_step.reward
    total_return += episode_return

  avg_return = total_return / num_episodes
  return avg_return.numpy()[0]


random_policy = random_tf_policy.RandomTFPolicy(train_env.time_step_spec(),
                                                train_env.action_spec())

compute_avg_return(eval_env, random_policy, num_eval_episodes)

# Please also see the metrics module for standard implementations of different
# metrics.
21.1

Coleção de dados

Como no tutorial DQN, configure o buffer de reprodução e a coleta de dados inicial com a política aleatória.

replay_buffer = tf_uniform_replay_buffer.TFUniformReplayBuffer(
    data_spec=agent.collect_data_spec,
    batch_size=train_env.batch_size,
    max_length=replay_buffer_capacity)

def collect_step(environment, policy):
  time_step = environment.current_time_step()
  action_step = policy.action(time_step)
  next_time_step = environment.step(action_step.action)
  traj = trajectory.from_transition(time_step, action_step, next_time_step)

  # Add trajectory to the replay buffer
  replay_buffer.add_batch(traj)

for _ in range(initial_collect_steps):
  collect_step(train_env, random_policy)

# This loop is so common in RL, that we provide standard implementations of
# these. For more details see the drivers module.

# Dataset generates trajectories with shape [BxTx...] where
# T = n_step_update + 1.
dataset = replay_buffer.as_dataset(
    num_parallel_calls=3, sample_batch_size=batch_size,
    num_steps=n_step_update + 1).prefetch(3)

iterator = iter(dataset)
WARNING:tensorflow:From /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/autograph/impl/api.py:382: ReplayBuffer.get_next (from tf_agents.replay_buffers.replay_buffer) is deprecated and will be removed in a future version.
Instructions for updating:
Use `as_dataset(..., single_deterministic_pass=False) instead.

Treinando o agente

O loop de treinamento envolve a coleta de dados do ambiente e a otimização das redes do agente. Ao longo do caminho, iremos ocasionalmente avaliar a política do agente para ver como estamos nos saindo.

O procedimento a seguir levará cerca de 7 minutos para ser executado.

try:
  %%time
except:
  pass

# (Optional) Optimize by wrapping some of the code in a graph using TF function.
agent.train = common.function(agent.train)

# Reset the train step
agent.train_step_counter.assign(0)

# Evaluate the agent's policy once before training.
avg_return = compute_avg_return(eval_env, agent.policy, num_eval_episodes)
returns = [avg_return]

for _ in range(num_iterations):

  # Collect a few steps using collect_policy and save to the replay buffer.
  for _ in range(collect_steps_per_iteration):
    collect_step(train_env, agent.collect_policy)

  # Sample a batch of data from the buffer and update the agent's network.
  experience, unused_info = next(iterator)
  train_loss = agent.train(experience)

  step = agent.train_step_counter.numpy()

  if step % log_interval == 0:
    print('step = {0}: loss = {1}'.format(step, train_loss.loss))

  if step % eval_interval == 0:
    avg_return = compute_avg_return(eval_env, agent.policy, num_eval_episodes)
    print('step = {0}: Average Return = {1:.2f}'.format(step, avg_return))
    returns.append(avg_return)
WARNING:tensorflow:From /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/util/dispatch.py:206: calling foldr_v2 (from tensorflow.python.ops.functional_ops) with back_prop=False is deprecated and will be removed in a future version.
Instructions for updating:
back_prop=False is deprecated. Consider using tf.stop_gradient instead.
Instead of:
results = tf.foldr(fn, elems, back_prop=False)
Use:
results = tf.nest.map_structure(tf.stop_gradient, tf.foldr(fn, elems))
step = 200: loss = 3.073254108428955
step = 400: loss = 1.7349590063095093
step = 600: loss = 1.792757272720337
step = 800: loss = 1.8498374223709106
step = 1000: loss = 1.5893661975860596
step = 1000: Average Return = 85.80
step = 1200: loss = 1.5345051288604736
step = 1400: loss = 1.0927748680114746
step = 1600: loss = 1.3448643684387207
step = 1800: loss = 1.3334734439849854
step = 2000: loss = 1.2656627893447876
step = 2000: Average Return = 104.50
step = 2200: loss = 0.9804798364639282
step = 2400: loss = 1.0801945924758911
step = 2600: loss = 0.8012468814849854
step = 2800: loss = 1.2241151332855225
step = 3000: loss = 1.101938009262085
step = 3000: Average Return = 75.00
step = 3200: loss = 1.1817928552627563
step = 3400: loss = 0.9651275873184204
step = 3600: loss = 1.1266820430755615
step = 3800: loss = 0.840592622756958
step = 4000: loss = 0.9282000064849854
step = 4000: Average Return = 129.00
step = 4200: loss = 1.0468955039978027
step = 4400: loss = 1.0415769815444946
step = 4600: loss = 1.0660417079925537
step = 4800: loss = 0.9771052002906799
step = 5000: loss = 0.8960349559783936
step = 5000: Average Return = 126.60
step = 5200: loss = 0.925206184387207
step = 5400: loss = 0.6512923240661621
step = 5600: loss = 0.7966318130493164
step = 5800: loss = 0.7415628433227539
step = 6000: loss = 0.897962212562561
step = 6000: Average Return = 109.10
step = 6200: loss = 0.8640482425689697
step = 6400: loss = 1.038527250289917
step = 6600: loss = 1.017322063446045
step = 6800: loss = 0.7789400815963745
step = 7000: loss = 0.9924593567848206
step = 7000: Average Return = 122.80
step = 7200: loss = 0.5298917293548584
step = 7400: loss = 0.5959913730621338
step = 7600: loss = 0.6432434320449829
step = 7800: loss = 0.6911174058914185
step = 8000: loss = 0.5797613263130188
step = 8000: Average Return = 144.50
step = 8200: loss = 0.5593736171722412
step = 8400: loss = 0.6327507495880127
step = 8600: loss = 0.5596221089363098
step = 8800: loss = 0.6492353677749634
step = 9000: loss = 0.6443781852722168
step = 9000: Average Return = 135.80
step = 9200: loss = 0.8031193614006042
step = 9400: loss = 0.541293740272522
step = 9600: loss = 0.643537700176239
step = 9800: loss = 0.7450239062309265
step = 10000: loss = 0.5937005281448364
step = 10000: Average Return = 202.80
step = 10200: loss = 0.4400593042373657
step = 10400: loss = 0.7001190185546875
step = 10600: loss = 0.56056809425354
step = 10800: loss = 0.7413918375968933
step = 11000: loss = 0.6136442422866821
step = 11000: Average Return = 126.10
step = 11200: loss = 0.6289737224578857
step = 11400: loss = 0.6546649932861328
step = 11600: loss = 0.4087551534175873
step = 11800: loss = 0.6893659830093384
step = 12000: loss = 0.5798210501670837
step = 12000: Average Return = 152.90
step = 12200: loss = 0.4694007635116577
step = 12400: loss = 0.6825675368309021
step = 12600: loss = 0.675622284412384
step = 12800: loss = 0.6773465275764465
step = 13000: loss = 0.35596776008605957
step = 13000: Average Return = 141.70
step = 13200: loss = 0.6894246339797974
step = 13400: loss = 0.7381303310394287
step = 13600: loss = 0.675399661064148
step = 13800: loss = 0.6918181777000427
step = 14000: loss = 0.5304032564163208
step = 14000: Average Return = 146.10
step = 14200: loss = 0.6360524892807007
step = 14400: loss = 0.5732524394989014
step = 14600: loss = 0.4850385785102844
step = 14800: loss = 0.37867382168769836
step = 15000: loss = 0.46798139810562134
step = 15000: Average Return = 154.10

Visualização

Enredos

Podemos traçar o retorno versus as etapas globais para ver o desempenho de nosso agente. No Cartpole-v1 , o ambiente oferece uma recompensa de +1 para cada vez que a barra permanece levantada e, como o número máximo de etapas é 500, o retorno máximo possível também é 500.

steps = range(0, num_iterations + 1, eval_interval)
plt.plot(steps, returns)
plt.ylabel('Average Return')
plt.xlabel('Step')
plt.ylim(top=550)
(6.240000247955322, 550.0)

png

Vídeos

É útil visualizar o desempenho de um agente, renderizando o ambiente em cada etapa. Antes de fazermos isso, vamos primeiro criar uma função para incorporar vídeos neste colab.

def embed_mp4(filename):
  """Embeds an mp4 file in the notebook."""
  video = open(filename,'rb').read()
  b64 = base64.b64encode(video)
  tag = '''
  <video width="640" height="480" controls>
    <source src="data:video/mp4;base64,{0}" type="video/mp4">
  Your browser does not support the video tag.
  </video>'''.format(b64.decode())

  return IPython.display.HTML(tag)

O código a seguir visualiza a política do agente para alguns episódios:

num_episodes = 3
video_filename = 'imageio.mp4'
with imageio.get_writer(video_filename, fps=60) as video:
  for _ in range(num_episodes):
    time_step = eval_env.reset()
    video.append_data(eval_py_env.render())
    while not time_step.is_last():
      action_step = agent.policy.action(time_step)
      time_step = eval_env.step(action_step.action)
      video.append_data(eval_py_env.render())

embed_mp4(video_filename)
WARNING:root:IMAGEIO FFMPEG_WRITER WARNING: input image is not divisible by macro_block_size=16, resizing from (400, 600) to (400, 608) to ensure video compatibility with most codecs and players. To prevent resizing, make your input image divisible by the macro_block_size or set the macro_block_size to None (risking incompatibility). You may also see a FFMPEG warning concerning speedloss due to data not being aligned.

O C51 tende a se sair um pouco melhor do que o DQN no CartPole-v1, mas a diferença entre os dois agentes se torna cada vez mais significativa em ambientes cada vez mais complexos. Por exemplo, no benchmark Atari 2600 completo, C51 demonstra uma melhoria de pontuação média de 126% sobre DQN após normalizar em relação a um agente aleatório. Melhorias adicionais podem ser obtidas incluindo atualizações de n etapas.

Para um mergulho mais profundo no algoritmo C51, consulte A Distributional Perspective on Reinforcement Learning (2017) .