<jupyter_start><jupyter_text>How to do it... Load in the required packages<jupyter_code>!pip install tf-agents

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_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
tf.compat.v1.reset_default_graph()
tf.compat.v1.enable_resource_variables()
tf.compat.v1.enable_v2_behavior()
nest = tf.compat.v2.nest

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<jupyter_output>Requirement already satisfied: tf-agents in /home/wil/anaconda3/lib/python3.7/site-packages (0.12.0)
Requirement already satisfied: protobuf>=3.11.3 in /home/wil/anaconda3/lib/python3.7/site-packages (from tf-agents) (3.19.4)
Requirement already satisfied: pillow in /home/wil/anaconda3/lib/python3.7/site-packages (from tf-agents) (9.0.1)
Requirement already satisfied: tensorflow-probability>=0.16.0 in /home/wil/anaconda3/lib/python3.7/site-packages (from tf-agents) (0.16.0)
Requirement already satisfied: gym>=0.17.0 in /home/wil/anaconda3/lib/python3.7/site-packages (from tf-agents) (0.22.0)
Requirement already satisfied: cloudpickle>=1.3 in /home/wil/anaconda3/lib/python3.7/site-packages (from tf-agents) (2.0.0)
Requirement already satisfied: typing-extensions>=3.7.4.3 in /home/wil/anaconda3/lib/python3.7/site-packages (from tf-agents) (4.1.1)
Requirement already satisfied: gin-config>=0.4.0 in /home/wil/anaconda3/lib/python3.7/site-packages (from tf-agents) (0.5.0)
Requirement alread[...]<jupyter_text>Define hyperparameters<jupyter_code>batch_size = 2

num_iterations = 100
steps_per_loop = 1<jupyter_output><empty_output><jupyter_text>Our first function is a context sampler used to generate observations coming from the environment.<jupyter_code>def context_sampling_fn(batch_size):
    
    def _context_sampling_fn():
        return np.random.randint(0, 2, [batch_size, 2]).astype(np.float32)
    return _context_sampling_fn<jupyter_output><empty_output><jupyter_text>We define a class for calculating the rewards.<jupyter_code>class CalculateReward(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)
        return (mu > 0 ) + 0<jupyter_output><empty_output><jupyter_text>We use this class to define the rewards per arm.<jupyter_code>arm0_param = [2, -1]
arm1_param = [1, -1]
arm2_param = [-1, 1]
arm3_param = [0, 0]

arm0_reward_fn = CalculateReward(arm0_param, 1)
arm1_reward_fn = CalculateReward(arm1_param, 1)
arm2_reward_fn = CalculateReward(arm2_param, 1)
arm3_reward_fn = CalculateReward(arm3_param, 1)<jupyter_output><empty_output><jupyter_text>And the final part of the function setup is to calculate the optimal rewards for a given context<jupyter_code>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)),
        tf.linalg.matvec(observation, tf.cast(arm3_param, dtype=tf.float32)),
    ]
    optimal_action_reward = tf.reduce_max(expected_reward_for_arms, axis=0)
    
    return optimal_action_reward<jupyter_output><empty_output><jupyter_text>We assume a stationary environment for this example<jupyter_code>environment = tf_py_environment.TFPyEnvironment(
    sspe.StationaryStochasticPyEnvironment(
        context_sampling_fn(batch_size),
        [arm0_reward_fn, arm1_reward_fn, arm2_reward_fn, arm3_reward_fn],
        batch_size=batch_size))<jupyter_output><empty_output><jupyter_text>And finally we can instantiate an agent that implements a bandit algorithm.<jupyter_code>observation_spec = tensor_spec.TensorSpec([2], 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)<jupyter_output><empty_output><jupyter_text>We calculate the regret which is the delta between the actual reward and expected reward<jupyter_code>regret_metric = tf_metrics.RegretMetric(compute_optimal_reward)<jupyter_output><empty_output><jupyter_text>And finally we can train our agent<jupyter_code>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())<jupyter_output>WARNING:tensorflow:From <ipython-input-10-f905f7342006>:18: ReplayBuffer.gather_all (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=True)` instead.<jupyter_text>It is possible to visualize our results as seen below...<jupyter_code>plt.plot(regret_values)
plt.ylabel('Average Regret')
plt.xlabel('Number of Iterations')<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.09648273636656986, 'fraction_numerical': 0.02452404001290739, 'mean_word_length': 5.857300884955753, 'pattern_counts': {'":': 0, '<': 44, '<?xml version=': 0, '>': 51, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 1, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '11', 'dir': 'jupyter-structured-clean-dedup', 'id': '26471', 'license': 'MIT', 'n_tokens_mistral': 2226, 'n_tokens_neox': 2125, 'n_words': 420, 'path': 'Chapter 11: Reinforcement Learning with TensorFlow and TF-Agents/Multi Armed Bandit.ipynb', 'repo_name': 'Young-Picasso/TensorFlow2.x_Cookbook', 'size': '6198'}