bandits_tutorial.ipynb

Kernel: Python 3

Copyright 2023 The TF-Agents Authors.

#@title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

Tutorial sobre bandidos multibrazo en TF-Agents

Introducción

Ver en TensorFlow.org

Ejecutar en Google Colab

Ver código fuente en GitHub

Descargar bloc de notas

Preparación

Si todavía no ha instalado las siguientes dependencias, ejecute estos comandos:

!pip install tf-agents

Importaciones

import abc
import numpy as np
import tensorflow as tf

from tf_agents.agents import tf_agent
from tf_agents.drivers import driver
from tf_agents.environments import py_environment
from tf_agents.environments import tf_environment
from tf_agents.environments import tf_py_environment
from tf_agents.policies import tf_policy
from tf_agents.specs import array_spec
from tf_agents.specs import tensor_spec
from tf_agents.trajectories import time_step as ts
from tf_agents.trajectories import trajectory
from tf_agents.trajectories import policy_step

nest = tf.nest

Introducción

El problema de los bandidos multibrazo (MAB, por sus siglas en inglés) es un caso especial de Aprendizaje por Refuerzo: un agente recopila recompensas en un entorno a través de la ejecución de ciertas acciones tras observar parte del estado del entorno. La principal diferencia entre el RL general y los MAB es que, en los MAB, asumimos que la acción ejecutada por el agente no influye en el siguiente estado del entorno. Por lo tanto, los agentes no modelan las transiciones de estado, no acreditan recompensas a las acciones pasadas ni "planifican con antelación" para llegar a estados con abundantes recompensas.

Como en otros dominios del RL, el objetivo de un agente del MAB es el de encontrar una política que recopile la mayor recompensa posible. Sin embargo, sería un error que siempre se intente explotar la acción que promete la mayor recompensa, porque entonces podríamos perder de vista mejores acciones por no explorar lo suficiente. Este es el problema más importante que debemos resolver en (MAB), a menudo conocido como la paradoja exploración-explotación.

Los entornos, las políticas y los agentes de bandidos para MAB se pueden encontrar en los subdirectorios de tf_agents/bandits.

Entornos

En TF-Agents, la clase de entorno cumple la función de dar información sobre el estado actual (esto se llama observación o contexto), recibir una acción como entrada, realizar una transición de estado y generar una recompensa. Esta clase también se encarga del reinicio cuando un episodio termina, para que un nuevo episodio pueda comenzar. Esto se consigue llamando a una función reset cuando un estado se etiqueta como "último" del episodio.

Para obtener más información, consulte el tutorial sobre entornos de TF-Agents.

Como se mencionó anteriormente, MAB es diferente del RL general en el sentido de que las acciones no influyen sobre la siguiente observación. Otra diferencia es que en Bandits no hay "episodios": cada paso de tiempo se inicia con una nueva observación, independientemente de los pasos de tiempo anteriores.

Para asegurarnos de que las observaciones sean independientes y para abstraer el concepto de episodios de RL, introducimos subclases de PyEnvironment y TFEnvironment: BanditPyEnvironment y BanditTFEnvironment. Estas clases exponen dos funciones miembro privadas que deben ser implementadas por el usuario:

@abc.abstractmethod
def _observe(self):

y

@abc.abstractmethod
def _apply_action(self, action):

La función _observe devuelve una observación. Luego, la política elige una acción en función de esta observación. _apply_action recibe esa acción como entrada y devuelve la recompensa correspondiente. Estas funciones miembro privadas son llamadas por las funciones reset y step, respectivamente.

class BanditPyEnvironment(py_environment.PyEnvironment):

  def __init__(self, observation_spec, action_spec):
    self._observation_spec = observation_spec
    self._action_spec = action_spec
    super(BanditPyEnvironment, self).__init__()

  # Helper functions.
  def action_spec(self):
    return self._action_spec

  def observation_spec(self):
    return self._observation_spec

  def _empty_observation(self):
    return tf.nest.map_structure(lambda x: np.zeros(x.shape, x.dtype),
                                 self.observation_spec())

  # These two functions below should not be overridden by subclasses.
  def _reset(self):
    """Returns a time step containing an observation."""
    return ts.restart(self._observe(), batch_size=self.batch_size)

  def _step(self, action):
    """Returns a time step containing the reward for the action taken."""
    reward = self._apply_action(action)
    return ts.termination(self._observe(), reward)

  # These two functions below are to be implemented in subclasses.
  @abc.abstractmethod
  def _observe(self):
    """Returns an observation."""

  @abc.abstractmethod
  def _apply_action(self, action):
    """Applies `action` to the Environment and returns the corresponding reward.
    """

La clase abstracta provisional que acabamos de describir implementa las funciones _reset y _step de PyEnvironment y expone las funciones abstractas _observe y _apply_action para que las implementen las subclases.

Una clase de entorno de ejemplo simple

La clase que se presenta a continuación ofrece un entorno muy simple en el que la observación es un número entero aleatorio comprendido entre -2 y 2, existen 3 acciones posibles (0, 1, 2) y la recompensa es el producto de la acción y la observación.

class SimplePyEnvironment(BanditPyEnvironment):

  def __init__(self):
    action_spec = array_spec.BoundedArraySpec(
        shape=(), dtype=np.int32, minimum=0, maximum=2, name='action')
    observation_spec = array_spec.BoundedArraySpec(
        shape=(1,), dtype=np.int32, minimum=-2, maximum=2, name='observation')
    super(SimplePyEnvironment, self).__init__(observation_spec, action_spec)

  def _observe(self):
    self._observation = np.random.randint(-2, 3, (1,), dtype='int32')
    return self._observation

  def _apply_action(self, action):
    return action * self._observation

Ahora podemos utilizar este entorno para obtener observaciones y recibir recompensas por nuestras acciones.

environment = SimplePyEnvironment()
observation = environment.reset().observation
print("observation: %d" % observation)

action = 2 #@param

print("action: %d" % action)
reward = environment.step(action).reward
print("reward: %f" % reward)

Entornos de TF

Se puede definir un entorno bandido al subclasificar BanditTFEnvironment o, de manera similar a los entornos de RL, se puede definir un BanditPyEnvironment y envolverlo con TFPyEnvironment. Para simplificar, en este tutorial optaremos por esta última opción.

tf_environment = tf_py_environment.TFPyEnvironment(environment)

Políticas

En un problema de bandidos, las políticas funcionan igual que en un problema de RL: ofrecen una acción (o una distribución de acciones), a partir de una observación que funciona como entrada.

Para obtener más información, consulte el tutorial sobre políticas de TF-Agents.

Del mismo modo que ocurre con los entornos, hay dos formas de construir una política: se puede crear una PyPolicy y envolverla con TFPyPolicy, o crear directamente una TFPy. Aquí optamos por el método directo.

Como este ejemplo es bastante sencillo, podemos definir la política óptima manualmente. La acción solo depende del signo de la observación, 0 cuando es negativa y 2 cuando es positiva.

class SignPolicy(tf_policy.TFPolicy):
  def __init__(self):
    observation_spec = tensor_spec.BoundedTensorSpec(
        shape=(1,), dtype=tf.int32, minimum=-2, maximum=2)
    time_step_spec = ts.time_step_spec(observation_spec)

    action_spec = tensor_spec.BoundedTensorSpec(
        shape=(), dtype=tf.int32, minimum=0, maximum=2)

    super(SignPolicy, self).__init__(time_step_spec=time_step_spec,
                                     action_spec=action_spec)
  def _distribution(self, time_step):
    pass

  def _variables(self):
    return ()

  def _action(self, time_step, policy_state, seed):
    observation_sign = tf.cast(tf.sign(time_step.observation[0]), dtype=tf.int32)
    action = observation_sign + 1
    return policy_step.PolicyStep(action, policy_state)

Ahora podemos solicitar una observación del entorno, llamar a la política para que elija una acción y, entonces, el entorno emitirá la recompensa:

sign_policy = SignPolicy()

current_time_step = tf_environment.reset()
print('Observation:')
print (current_time_step.observation)
action = sign_policy.action(current_time_step).action
print('Action:')
print (action)
reward = tf_environment.step(action).reward
print('Reward:')
print(reward)

La forma en que se implementan los entornos bandidos garantiza que cada vez que damos un paso, recibimos tanto la recompensa por la acción efectuada como la siguiente observación.

step = tf_environment.reset()
action = 1
next_step = tf_environment.step(action)
reward = next_step.reward
next_observation = next_step.observation
print("Reward: ")
print(reward)
print("Next observation:")
print(next_observation)

Agentes

Ahora que ya explicamos los entornos bandidos y las políticas bandidas, debemos definir también los agentes bandidos, que se encargan de cambiar la política en función de muestras de entrenamiento.

La API para los agentes bandido es similar a la de los agentes de RL: el agente solo tiene que implementar los métodos _initialize y _train y definir una policy y una collect_policy.

Un entorno más complicado

Antes de que escribamos nuestro agente bandido, necesitamos contar con un entorno que sea un poco más difícil de entender. Para hacerlo un poco más interesante, el siguiente entorno siempre dará reward = observation * action o reward = -observation * action. Esto se decidirá cuando se inicialice el entorno.

class TwoWayPyEnvironment(BanditPyEnvironment):

  def __init__(self):
    action_spec = array_spec.BoundedArraySpec(
        shape=(), dtype=np.int32, minimum=0, maximum=2, name='action')
    observation_spec = array_spec.BoundedArraySpec(
        shape=(1,), dtype=np.int32, minimum=-2, maximum=2, name='observation')

    # Flipping the sign with probability 1/2.
    self._reward_sign = 2 * np.random.randint(2) - 1
    print("reward sign:")
    print(self._reward_sign)

    super(TwoWayPyEnvironment, self).__init__(observation_spec, action_spec)

  def _observe(self):
    self._observation = np.random.randint(-2, 3, (1,), dtype='int32')
    return self._observation

  def _apply_action(self, action):
    return self._reward_sign * action * self._observation[0]

two_way_tf_environment = tf_py_environment.TFPyEnvironment(TwoWayPyEnvironment())

Una política más complicada

Un entorno más complicado requiere una política más complicada. Necesitamos una política que detecte el comportamiento del entorno subyacente. Hay tres situaciones que la política debe gestionar:

1. El agente aún no ha detectado qué versión del entorno se está ejecutando.

2. El agente detectó que se está ejecutando la versión original del entorno.

3. El agente detectó que se está ejecutando la versión invertida del entorno.

Definimos una tf_variable nombrada _situation para almacenar esta información codificada como valores en [0, 2], luego hacemos que la política se comporte en consecuencia.

class TwoWaySignPolicy(tf_policy.TFPolicy):
  def __init__(self, situation):
    observation_spec = tensor_spec.BoundedTensorSpec(
        shape=(1,), dtype=tf.int32, minimum=-2, maximum=2)
    action_spec = tensor_spec.BoundedTensorSpec(
        shape=(), dtype=tf.int32, minimum=0, maximum=2)
    time_step_spec = ts.time_step_spec(observation_spec)
    self._situation = situation
    super(TwoWaySignPolicy, self).__init__(time_step_spec=time_step_spec,
                                           action_spec=action_spec)
  def _distribution(self, time_step):
    pass

  def _variables(self):
    return [self._situation]

  def _action(self, time_step, policy_state, seed):
    sign = tf.cast(tf.sign(time_step.observation[0, 0]), dtype=tf.int32)
    def case_unknown_fn():
      # Choose 1 so that we get information on the sign.
      return tf.constant(1, shape=(1,))

    # Choose 0 or 2, depending on the situation and the sign of the observation.
    def case_normal_fn():
      return tf.constant(sign + 1, shape=(1,))
    def case_flipped_fn():
      return tf.constant(1 - sign, shape=(1,))

    cases = [(tf.equal(self._situation, 0), case_unknown_fn),
             (tf.equal(self._situation, 1), case_normal_fn),
             (tf.equal(self._situation, 2), case_flipped_fn)]
    action = tf.case(cases, exclusive=True)
    return policy_step.PolicyStep(action, policy_state)

El agente

Ahora debemos definir el agente que detecta el signo del entorno y establece la política adecuada.

class SignAgent(tf_agent.TFAgent):
  def __init__(self):
    self._situation = tf.Variable(0, dtype=tf.int32)
    policy = TwoWaySignPolicy(self._situation)
    time_step_spec = policy.time_step_spec
    action_spec = policy.action_spec
    super(SignAgent, self).__init__(time_step_spec=time_step_spec,
                                    action_spec=action_spec,
                                    policy=policy,
                                    collect_policy=policy,
                                    train_sequence_length=None)

  def _initialize(self):
    return tf.compat.v1.variables_initializer(self.variables)

  def _train(self, experience, weights=None):
    observation = experience.observation
    action = experience.action
    reward = experience.reward

    # We only need to change the value of the situation variable if it is
    # unknown (0) right now, and we can infer the situation only if the
    # observation is not 0.
    needs_action = tf.logical_and(tf.equal(self._situation, 0),
                                  tf.not_equal(reward, 0))


    def new_situation_fn():
      """This returns either 1 or 2, depending on the signs."""
      return (3 - tf.sign(tf.cast(observation[0, 0, 0], dtype=tf.int32) *
                          tf.cast(action[0, 0], dtype=tf.int32) *
                          tf.cast(reward[0, 0], dtype=tf.int32))) / 2

    new_situation = tf.cond(needs_action,
                            new_situation_fn,
                            lambda: self._situation)
    new_situation = tf.cast(new_situation, tf.int32)
    tf.compat.v1.assign(self._situation, new_situation)
    return tf_agent.LossInfo((), ())

sign_agent = SignAgent()

En el código anterior, el agente define la política y el agente y la política comparten la variable situation.

Además, el parámetro experience de la función _train corresponde a una trayectoria:

Trayectorias

En TF-Agents, trajectories son tuplas nombradas que contienen muestras de los pasos anteriores. Luego, los agentes usan estas muestras para entrenar y actualizar la política. En el RL, las trayectorias deben contener información sobre el estado actual, el siguiente estado, y sobre si el episodio actual ha terminado. Como en el mundo bandido no necesitamos estos datos, configuramos una función ayudante para crear una trayectoria:

# We need to add another dimension here because the agent expects the
# trajectory of shape [batch_size, time, ...], but in this tutorial we assume
# that both batch size and time are 1. Hence all the expand_dims.

def trajectory_for_bandit(initial_step, action_step, final_step):
  return trajectory.Trajectory(observation=tf.expand_dims(initial_step.observation, 0),
                               action=tf.expand_dims(action_step.action, 0),
                               policy_info=action_step.info,
                               reward=tf.expand_dims(final_step.reward, 0),
                               discount=tf.expand_dims(final_step.discount, 0),
                               step_type=tf.expand_dims(initial_step.step_type, 0),
                               next_step_type=tf.expand_dims(final_step.step_type, 0))

Cómo entrenar un agente

Ya tenemos todo lo que necesitamos para entrenar nuestro agente bandido.

step = two_way_tf_environment.reset()
for _ in range(10):
  action_step = sign_agent.collect_policy.action(step)
  next_step = two_way_tf_environment.step(action_step.action)
  experience = trajectory_for_bandit(step, action_step, next_step)
  print(experience)
  sign_agent.train(experience)
  step = next_step

A partir de la salida se puede ver que después del segundo paso (a menos que la observación fuera 0 en el primer paso), la política elige la acción de forma correcta y por lo tanto la recompensa obtenida es siempre no negativa.

Un ejemplo real de bandido contextual

En lo que resta de este tutorial, usamos los entornos y los agentes preimplementados de la biblioteca TF-Agents Bandits.

# Imports for example.
from tf_agents.bandits.agents import lin_ucb_agent
from tf_agents.bandits.environments import stationary_stochastic_py_environment as sspe
from tf_agents.bandits.metrics import tf_metrics
from tf_agents.drivers import dynamic_step_driver
from tf_agents.replay_buffers import tf_uniform_replay_buffer

import matplotlib.pyplot as plt

Entorno estocástico estacionario con funciones de recompensa lineales

El entorno que se usa en este ejemplo es StationaryStochasticPyEnvironment. Este entorno toma como parámetro una función (generalmente ruidosa) para ofrecer observaciones (contexto), y para cada brazo toma una función (también ruidosa) que calcula la recompensa en función de la observación dada. En nuestro ejemplo, tomamos muestras del contexto de un cubo d-dimensional uniformemente, y las funciones de recompensa son funciones lineales del contexto, más un poco de ruido gaussiano.

batch_size = 2 # @param
arm0_param = [-3, 0, 1, -2] # @param
arm1_param = [1, -2, 3, 0] # @param
arm2_param = [0, 0, 1, 1] # @param
def context_sampling_fn(batch_size):
  """Contexts from [-10, 10]^4."""
  def _context_sampling_fn():
    return np.random.randint(-10, 10, [batch_size, 4]).astype(np.float32)
  return _context_sampling_fn

class LinearNormalReward(object):
  """A class that acts as linear reward function when called."""
  def __init__(self, theta, sigma):
    self.theta = theta
    self.sigma = sigma
  def __call__(self, x):
    mu = np.dot(x, self.theta)
    return np.random.normal(mu, self.sigma)

arm0_reward_fn = LinearNormalReward(arm0_param, 1)
arm1_reward_fn = LinearNormalReward(arm1_param, 1)
arm2_reward_fn = LinearNormalReward(arm2_param, 1)

environment = tf_py_environment.TFPyEnvironment(
    sspe.StationaryStochasticPyEnvironment(
        context_sampling_fn(batch_size),
        [arm0_reward_fn, arm1_reward_fn, arm2_reward_fn],
        batch_size=batch_size))

El agente LinUCB

El agente que se muestra a continuación implementa el algoritmo LinUCB.

observation_spec = tensor_spec.TensorSpec([4], tf.float32)
time_step_spec = ts.time_step_spec(observation_spec)
action_spec = tensor_spec.BoundedTensorSpec(
    dtype=tf.int32, shape=(), minimum=0, maximum=2)

agent = lin_ucb_agent.LinearUCBAgent(time_step_spec=time_step_spec,
                                     action_spec=action_spec)

Métrica de arrepentimiento

La métrica más importante de Bandits es arrepentimiento, y se calcula como la diferencia entre la recompensa obtenida por el agente y la recompensa esperada por una política de oráculo que tenga acceso a las funciones de recompensa del entorno. Por lo tanto, RegretMetric necesita una función baseline_reward_fn que calcule la mejor recompensa esperada alcanzable a partir de una observación. Para nuestro ejemplo, tenemos que tomar el máximo de los equivalentes sin ruido de las funciones de recompensa que ya hemos definido para el entorno.

def compute_optimal_reward(observation):
  expected_reward_for_arms = [
      tf.linalg.matvec(observation, tf.cast(arm0_param, dtype=tf.float32)),
      tf.linalg.matvec(observation, tf.cast(arm1_param, dtype=tf.float32)),
      tf.linalg.matvec(observation, tf.cast(arm2_param, dtype=tf.float32))]
  optimal_action_reward = tf.reduce_max(expected_reward_for_arms, axis=0)
  return optimal_action_reward

regret_metric = tf_metrics.RegretMetric(compute_optimal_reward)

Entrenamiento

En este paso, reunimos todos los componentes que presentamos anteriormente: el entorno, la política y el agente. Ejecutamos la política sobre el entorno y generamos datos de entrenamiento con la ayuda de un controlador, y capacitamos al agente con los datos.

Tenga en cuenta que hay dos parámetros que se unen para especificar el número de pasos dados. num_iterations especifica cuántas veces ejecutamos el bucle entrenador, mientras que el controlador dará steps_per_loop pasos por iteración. La principal razón para mantener ambos parámetros en algunas de estas operaciones es que algunas operaciones se ejecutan por iteración, mientras que a otras las ejecuta en cada paso el controlador. Por ejemplo, la función train del agente se llama una sola vez por iteración. La diferencia radica en que, si entrenamos más a menudo, nuestra política será más "fresca"; por otro lado, entrenar en lotes más grandes puede resultar más eficiente en función del tiempo.

num_iterations = 90 # @param
steps_per_loop = 1 # @param

replay_buffer = tf_uniform_replay_buffer.TFUniformReplayBuffer(
    data_spec=agent.policy.trajectory_spec,
    batch_size=batch_size,
    max_length=steps_per_loop)

observers = [replay_buffer.add_batch, regret_metric]

driver = dynamic_step_driver.DynamicStepDriver(
    env=environment,
    policy=agent.collect_policy,
    num_steps=steps_per_loop * batch_size,
    observers=observers)

regret_values = []

for _ in range(num_iterations):
  driver.run()
  loss_info = agent.train(replay_buffer.gather_all())
  replay_buffer.clear()
  regret_values.append(regret_metric.result())

plt.plot(regret_values)
plt.ylabel('Average Regret')
plt.xlabel('Number of Iterations')

Tras ejecutar el último fragmento de código, el gráfico resultante (con suerte) mostrará que el arrepentimiento medio disminuye a medida que el agente se entrena y la política mejora su capacidad de predecir cuál es la acción correcta, a partir de la observación.

Siguientes pasos

Para ver más ejemplos prácticos, consulte bandits/agents/examples que tiene ejemplos listos para ejecutar para diferentes agentes y entornos.

La biblioteca TF-Agents también es capaz de gestionar bandidos multibrazo con características por brazo. Para ello, sugerimos que se consulte el tutorial de bandido por brazo.