<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 Applying Contextual Bandits for Recommendation systems using Tensorflow and Cloud Storage
Learning objectives
Install and import required libraries.
Initialize and configure the MovieLens Environment.
Initialize the Agent.
Define and link the evaluation metrics.
Initialize & configure the Replay Buffer.
Setup and Train the model.
Inference with trained model & Tensorboard Evaluation.
Introduction
Multi-Armed Bandit (MAB) is a Machine Learning framework in which an agent has to select actions (arms) in order to maximize its cumulative reward in the long term. In each round, the agent receives some information about the current state (context), then it chooses an action based on this information and the experience gathered in previous rounds. At the end of each round, the agent receives the reward assiociated with the chosen action.
https
Step1:  2. Initializing and configuring the MovieLens Environment
Firstly we need to load the movielens.data csv file stored in cloud storage, load it locally and initilialze the MovielensPyenvironment with it. Refer here for guidance on it.
An environment in the TF-Agents Bandits library is a class that provides observations and reports rewards based on observations and actions.
We will be using the MovieLens environment. This environment implements the MovieLens 100K dataset, available at
Step2:  3. Initializing the Agent
Now that we have the environment query we reach the part where we define and initialize our policy and the Agent which will be our utilize that policy to make decisions given an observation. We have several policies
Step3:  4. Define and link the evaluation metrics
Just like you have metrics like accuracy/recall in supervised learning, in bandits we use the regret metric per episode. To calculate the regret, we need to know what the highest possible expected reward is in every time step. For that, we define the optimal_reward_fn.
Another similar metric is the number of times a suboptimal action was chosen. That requires the definition if the optimal_action_fn.
Step4:  5. Initialize & configure the Replay Buffer
Reinforcement learning algorithms use replay buffers to store trajectories of experience when executing a policy in an environment. During training, replay buffers are queried for a subset of the trajectories (either a sequential subset or a sample) to "replay" the agent's experience. Sampling from the replay buffer facilitate data re-use and breaks harmful co-relation between sequential data in RL, although in contextual bandits this isn't absolutely required but still helpful.
The replay buffer exposes several functions which allow you to manipulate the replay buffer in several ways. Read more on them [here] (https
Step5:  Now we have a Replay buffer but we also need something to fill it with. Often a common practice is to have 
the agent Interact with and collect experience with the environment, without actually learning from it ( i.e. only forward pass). This loop can  be either by you manually as shown here or you can do it using the DynamicStepDriver.
The data encountered by the driver at each step is saved in a named tuple called Trajectory and broadcast to a set of observers such as replay buffers and metrics. 
This Trajectory includes the observation from the environment, the action recommended by the policy, the reward obtained, the type of the current and the next step, etc. 
In order for the driver to fill the replay buffer with data, as well as to compute ongoing metrics, it needs acess to the add_batch, functionality of the buffer, and the metrics ( both step and regular). Refer here for more information aand example code on how initialize a step driver with observers.
Step7:  6. Setup and Train the Model
Here we provide you a helper function in order to save your agent, the metrics and its lighter policy seperately, while training the model. We make all the aspects into trackable objects and then use checkpoint to save as well warm restart a previous training. For more information on checkpoints and policy savers ( which will be used in the training loop below) refer here
Step8:  Now we have all the components ready to start training the model. Here is the process for Training the model
1. We first use the DynamicStepdriver instance to collect experience( trajectories) from the environment and fill up the replay buffer.
2. We then extract all the stored experience from the replay buffer by specfiying the batch size and num_steps the same as we initialized the driver with. We extract it as tf.dataset instance.
3. We then iterate on the tf.dataset and the first sample we draw actually has all the data batch_size*num_time_steps
4. the agent then trains on the acquired experience
5. the replay buffer is cleared to make space for new data
6. Log the metrics and store them on disk
7. Save the Agent ( via checkpoints) as well as the policy
We recommend doing the training for 15,000 loops with 2 steps per loop, and an agent alpha of 10.0
Step9:  Note
Step10:  One last task before starting the training
Step11:  <img src='./assets/example_tensorboard.png'>
7. Inferencing with trained model & Tensorboard Evaluation
Now that our model is trained, what if we want to determine which action to take given a new "context"
<ASSISTANT_TASK:>
Python Code:
!pip install --quiet --upgrade --force-reinstall tensorflow==2.4 tensorflow_probability==0.12.1 tensorflow-io==0.17.0 --use-feature=2020-resolver
!pip install tf_agents==0.7.1 --quiet gast==0.3.3 --upgrade --use-feature=2020-resolver
import functools
import os
from absl import app
from absl import flags
import tensorflow as tf  # pylint: disable=g-explicit-tensorflow-version-import
from tf_agents.bandits.agents import dropout_thompson_sampling_agent as dropout_ts_agent
from tf_agents.bandits.agents import lin_ucb_agent
from tf_agents.bandits.agents import linear_thompson_sampling_agent as lin_ts_agent
from tf_agents.bandits.agents import neural_epsilon_greedy_agent as eps_greedy_agent
from tf_agents.bandits.agents.examples.v2 import trainer
from tf_agents.bandits.environments import environment_utilities
#from tf_agents.bandits.environments import movielens_per_arm_py_environment
from tf_agents.bandits.environments import movielens_py_environment
from tf_agents.metrics import tf_metrics
from tf_agents.bandits.metrics import tf_metrics as tf_bandit_metrics
from tf_agents.bandits.networks import global_and_arm_feature_network
from tf_agents.environments import tf_py_environment
from tf_agents.networks import q_network
from tf_agents.drivers import dynamic_step_driver
from tf_agents.eval import metric_utils
from tf_agents.policies import policy_saver
from tf_agents.replay_buffers import tf_uniform_replay_buffer
from tf_agents.trajectories import time_step as ts
# If there are version / incompatibility errors, make sure you restarted the kernel and use !pip freeze in a new cell to check whether the correct TF and tf_agents version had been installed.
# Create target Directory if don't exist
from datetime import date
today = date.today()
fdate = date.today().strftime('%d_%m_%Y')
root_path = os.getcwd()
log_path = "{}/{}".format(root_path, fdate)
if not os.path.exists(log_path):
    os.mkdir(log_path)
    print("Directory {} Created".format(fdate))
else:    
    print("Directory {} already exists".format(fdate))
print("Full path is {}".format(log_path))
Explanation: Applying Contextual Bandits for Recommendation systems using Tensorflow and Cloud Storage
Learning objectives
Install and import required libraries.
Initialize and configure the MovieLens Environment.
Initialize the Agent.
Define and link the evaluation metrics.
Initialize & configure the Replay Buffer.
Setup and Train the model.
Inference with trained model & Tensorboard Evaluation.
Introduction
Multi-Armed Bandit (MAB) is a Machine Learning framework in which an agent has to select actions (arms) in order to maximize its cumulative reward in the long term. In each round, the agent receives some information about the current state (context), then it chooses an action based on this information and the experience gathered in previous rounds. At the end of each round, the agent receives the reward assiociated with the chosen action.
https://www.tensorflow.org/agents/tutorials/intro_bandit#multi-armed_bandits_and_reinforcement_learning
Each learning objective will correspond to a #TODO in this student lab notebook -- try to complete this notebook first and then review the solution notebook
1. Initial Setup: installing and importing required Libraries
End of explanation
# initialize the movielens pyenvironment with default parameters
NUM_ACTIONS = None # take this as 20
RANK_K = None # take rank as 20
BATCH_SIZE = None # take batch size as 8
data_path = "gs://ta-reinforecement-learning/dataset/movielens.data" # specify the path to the movielens.data OR get it from the GCS bucket
#TODO: replace the data path if needed
env = movielens_py_environment.MovieLensPyEnvironment(
        data_path, RANK_K, BATCH_SIZE, num_movies=NUM_ACTIONS)
environment = tf_py_environment.TFPyEnvironment(env)
Explanation: 2. Initializing and configuring the MovieLens Environment
Firstly we need to load the movielens.data csv file stored in cloud storage, load it locally and initilialze the MovielensPyenvironment with it. Refer here for guidance on it.
An environment in the TF-Agents Bandits library is a class that provides observations and reports rewards based on observations and actions.
We will be using the MovieLens environment. This environment implements the MovieLens 100K dataset, available at:
  https://www.kaggle.com/prajitdatta/movielens-100k-dataset
This dataset contains 100K ratings from m=943 users on n=1682 items. The ratings can be organized as a matrix A of size m-by-n.
Note that the ratings matrix is a <b>sparse matrix</b> i.e., only a subset of certain (user, movie) pairs is provided, since not all users have seen all movies.
In order for the environment to be able to compute a reasonable estimate of the reward, which represents how much a user i would enjoy a movie j,
the environment computes a dense approximation to this sparse matrix A.
In collaborative filtering, it is common practice to obtain this dense approximation by means of a low-rank matrix factorization of the matrix A.
The MovieLens environment uses truncated Singular Value Decomposition (SVD) (but other matrix factorization techniques could be potentially also used).
With truncated SVD of rank k, the matrix A is factorized as follows:
$A_k = U_k \Sigma_k V_k^T$,
where:
<li>$U_k$ is a matrix of orthogonal columns of size $m$-by-$k$,<\li>
<li>$V_k$ is a matrix of orthogonal columns of size $n$-by-$k$</li>
<li> $\Sigma_k$ is a diagonal matrix of size $k$-by-$k$ that holds the $k$ largest singular values of A.</li>
By splitting $\Sigma$ into $\sqrt{\Sigma_k} \sqrt{\Sigma_k}$, we can finally approximate the matrix A as a 
product of two factors $\tilde{U}$ and $\tilde{V}$ i.e.,
$A ~= \tilde{U} \tilde{V}^T$,
where $\tilde{U} = U_k \sqrt{\Sigma_k}$ and $\tilde{V} = V_k \sqrt{\Sigma_k}$
Once the matrix factorization has been computed, the environment caches it and uses it to compute the reward for recommending an movie `j` to a user `i` 
by retrieving the (`i`, `j`)-entry of matrix $A$.
Apart from computing the reward when the agent recommends a certain movie to a user, the environment is also responsible for generating observations that are given as input to the agent in order to make an informed decision. In order to generate a random observation, the environment samples a random row `i` from the matrix $\tilde{U}$. Once the agent selects movie `j` then the environment responds with the (`i`, `j`)-entry of matrix $A$.
End of explanation
# Replace these values by reading the above instructions carefully
EPSILON = 0
LAYERS = None
LR = 0
DROPOUT_RATE = 0
# Initialize the Qnetwork
network = q_network.QNetwork(
          input_tensor_spec=environment.time_step_spec().observation,
          action_spec=environment.action_spec(),
          fc_layer_params=LAYERS)
# Creating a neuron Epsilon greedy agent with an optimizer, 
# Epsilon exploration value, learning & dropout rate
# Replace all the `None` values with the required values
agent = eps_greedy_agent.NeuralEpsilonGreedyAgent(
  time_step_spec=None,# get the spec/format of the environment
  action_spec=None, # get the spec/format of the environment
  reward_network=None, #q network goes here
  optimizer=None #start w/ adam optimizer with a learning rate of .002
  epsilon=EPSILON) # we recommend an exploration of value of 1%)
Explanation: 3. Initializing the Agent
Now that we have the environment query we reach the part where we define and initialize our policy and the Agent which will be our utilize that policy to make decisions given an observation. We have several policies: as shown here:
NeuralEpsilonGreedyAgent: The neural episilon greed algorithm makes a value estimate for all the arms, and then chooses the best arm with the probaility (1-episilon) and any of the random arms with a probability of epsilon. this balances the exploration-exploitation tradeoff and epsilon is set to a small value like 10%. Example: In this example we have seven arms: one of each of the classes, and if we set episilon to say 10%, then 90% of the times the agent will choose the arm with the highest value estimate ( expplotiing the one most likely to be the predicted class) and 10% of the time it will choose a random arm from all of the 7 arms( thus exploring the other possibilities). Refer here for more information of the tensorflow agents version of the same.
Each Agent is initializied with a policy: which is essentially the function approximator ( be it linear or non linear) for estimating the Q values. Ther agen trains this policy, and the policy adds the exploration-exploitation component on top of this, and also chooses the action. In this example we will use a Deep Q Network as our value function, and we use the epsilon greedy on topof this to select the actions. In this case the action space would be 20 for 20 movies, the contextual state vector would be the dense user vector from the matrix decomposition. In applied situations, a dictionary mapping could be made from a user id to its dense representation to make it more convinient for the end user.
Step 1. Initialize the  Qnetwork, which takes in the state and returns the value function for each action. Define the Fully connected layer parameters to be (50, 50, 50) from the left to the right respectively.
Step 2. Creating a neuron Epsilon greedy agent  with an Adam Optimizer with Epsilon exploration value of 0.05, learning rate = 0.005, Dropout rate = 0.2. Feel free to experiment with these later to gauge their impact on the training later
Click here for reference on code example of how to create a  Q network and DQN Agents
End of explanation
# Making functions for computing optimal reward/action and attaching the env variable to it using partial functions, so it doesnt need to be passed with every invocation
optimal_reward_fn = functools.partial(
      environment_utilities.compute_optimal_reward_with_movielens_environment,
      environment=environment)
optimal_action_fn = functools.partial(
      environment_utilities.compute_optimal_action_with_movielens_environment,
      environment=environment)
# Initilializing the regret and suboptimal arms metric using the optimal reward and action functions
regret_metric = tf_bandit_metrics.RegretMetric(optimal_reward_fn)
suboptimal_arms_metric = tf_bandit_metrics.SuboptimalArmsMetric(
      optimal_action_fn)
step_metric = tf_metrics.EnvironmentSteps()
metrics = [tf_metrics.NumberOfEpisodes(),  #equivalent to number of steps in bandits problem
           regret_metric,  # measures regret
           suboptimal_arms_metric,  # number of times the suboptimal arms are pulled
           tf_metrics.AverageReturnMetric(batch_size=environment.batch_size)  # the average return
           ]
Explanation: 4. Define and link the evaluation metrics
Just like you have metrics like accuracy/recall in supervised learning, in bandits we use the regret metric per episode. To calculate the regret, we need to know what the highest possible expected reward is in every time step. For that, we define the optimal_reward_fn.
Another similar metric is the number of times a suboptimal action was chosen. That requires the definition if the optimal_action_fn.
End of explanation
#TODO
STEPS_PER_LOOP = None
# TFUniformReplayBuffer is the most commonly used replay buffer in TF-Agents. Use 'tf_uniform_replay_buffer.TFUniformReplayBuffer' to create one.
buf = None
Explanation: 5. Initialize & configure the Replay Buffer
Reinforcement learning algorithms use replay buffers to store trajectories of experience when executing a policy in an environment. During training, replay buffers are queried for a subset of the trajectories (either a sequential subset or a sample) to "replay" the agent's experience. Sampling from the replay buffer facilitate data re-use and breaks harmful co-relation between sequential data in RL, although in contextual bandits this isn't absolutely required but still helpful.
The replay buffer exposes several functions which allow you to manipulate the replay buffer in several ways. Read more on them [here] (https://www.tensorflow.org/agents/tutorials/5_replay_buffers_tutorial)
In this demo we would be using the TFUniformReplayBuffer for which we need to initialize the buffer spec with the spec of the trajectory of the agent's policy, a chosen batch size( number of trajectories to store), and the maximum length of the trajectory. ( this is the amount of sequential time steps which will be considered as one data point). so a batch of 3 with 2 time steps each would result in a tensor of shape (3,2). Since unlike regular RL problems, Contextual bandits have only one time step we can keep max_length =1, however since this tutorial is to enable you for RL problems as well, let set it to 2. Do not worry, any contextual bandit agent will internally
split the time steps inside each data point such that the effective batch size ends up being (6,1). 
Create a Tensorflow based UniformReplayBuffer And initialize it with an appropriate values.
Recommended:
    Batch size = 8
    Max length = 2 ( 2 time steps per item)
End of explanation
#TODO: setup the replay observer as a list to capture both metrics, step metrics and provide access to the function to load data from the driver into the buffer
replay_observer = None
driver = dynamic_step_driver.DynamicStepDriver(
      env=None,
      policy=None,
      num_steps=STEPS_PER_LOOP * environment.batch_size,
      observers=None)
Explanation: Now we have a Replay buffer but we also need something to fill it with. Often a common practice is to have 
the agent Interact with and collect experience with the environment, without actually learning from it ( i.e. only forward pass). This loop can  be either by you manually as shown here or you can do it using the DynamicStepDriver.
The data encountered by the driver at each step is saved in a named tuple called Trajectory and broadcast to a set of observers such as replay buffers and metrics. 
This Trajectory includes the observation from the environment, the action recommended by the policy, the reward obtained, the type of the current and the next step, etc. 
In order for the driver to fill the replay buffer with data, as well as to compute ongoing metrics, it needs acess to the add_batch, functionality of the buffer, and the metrics ( both step and regular). Refer here for more information aand example code on how initialize a step driver with observers.
End of explanation
AGENT_CHECKPOINT_NAME = 'agent'
STEP_CHECKPOINT_NAME = 'step'
CHECKPOINT_FILE_PREFIX = 'ckpt'
def restore_and_get_checkpoint_manager(root_dir, agent, metrics, step_metric):
    Restores from `root_dir` and returns a function that writes checkpoints.
    trackable_objects = {metric.name: metric for metric in metrics}
    trackable_objects[AGENT_CHECKPOINT_NAME] = agent
    trackable_objects[STEP_CHECKPOINT_NAME] = step_metric
    checkpoint = tf.train.Checkpoint(**trackable_objects)
    checkpoint_manager = tf.train.CheckpointManager(checkpoint=checkpoint,
                                                  directory=root_dir,
                                                  max_to_keep=5)
    latest = checkpoint_manager.latest_checkpoint
    if latest is not None:
        print('Restoring checkpoint from %s.', latest)
        checkpoint.restore(latest)
        print('Successfully restored to step %s.', step_metric.result())
    else:
        print('Did not find a pre-existing checkpoint. '
                 'Starting from scratch.')
    return checkpoint_manager
checkpoint_manager = restore_and_get_checkpoint_manager(
  log_path, agent, metrics, step_metric)
saver = policy_saver.PolicySaver(agent.policy)
summary_writer = tf.summary.create_file_writer(log_path)
summary_writer.set_as_default()
Explanation: 6. Setup and Train the Model
Here we provide you a helper function in order to save your agent, the metrics and its lighter policy seperately, while training the model. We make all the aspects into trackable objects and then use checkpoint to save as well warm restart a previous training. For more information on checkpoints and policy savers ( which will be used in the training loop below) refer here
End of explanation
#TODO 
# Replace `None` with the above given values
AGENT_ALPHA = None
TRAINING_LOOPS = None
Explanation: Now we have all the components ready to start training the model. Here is the process for Training the model
1. We first use the DynamicStepdriver instance to collect experience( trajectories) from the environment and fill up the replay buffer.
2. We then extract all the stored experience from the replay buffer by specfiying the batch size and num_steps the same as we initialized the driver with. We extract it as tf.dataset instance.
3. We then iterate on the tf.dataset and the first sample we draw actually has all the data batch_size*num_time_steps
4. the agent then trains on the acquired experience
5. the replay buffer is cleared to make space for new data
6. Log the metrics and store them on disk
7. Save the Agent ( via checkpoints) as well as the policy
We recommend doing the training for 15,000 loops with 2 steps per loop, and an agent alpha of 10.0
End of explanation
## TRAINING
#TOFINISH: define number of training loops and write the training function
TRAINING_LOOPS = None # We recommend doing 15k loops
import warnings
warnings.filterwarnings('ignore')
for _ in range(TRAINING_LOOPS):
    # step 1: We first use the DynamicStepdriver instance to collect experience
    #(trajectories) from the environment and fill up the replay buffer.
    
 
    # step 2: We then extract all the stored experience from the replay buffer by
    #specfiying the batch size and num_steps the same as we initialized the driver with.
    # We extract it as tf.dataset instance.
    
   
    # step 3: We then iterate on the tf.dataset and the first sample we draw 
    #actually has all the data batch_size*num_time_steps
    
    # step 4: The agent then trains on the acquired experience
    train_loss = agent.train(experience).loss
    
    # step 5:  the replay buffer is cleared to make space for new data
    
    
    # step 6: Log the metrics and store them on disk
    metric_utils.log_metrics(metrics)
    for metric in metrics:
        metric.tf_summaries(train_step=step_metric.result())
    
    # step 7: Save the Agent ( via checkpoints) as well as the policy
    checkpoint_manager.save()
    saver.save(os.path.join("./", 'policy_%d' % step_metric.result()))
Explanation: Note: The training will take around 50 minutes to complete and all the data are stored in the log_path directory.
End of explanation
print("tensorboard dev upload --logdir {} --name \"(optional) My latest experiment\" --description \"(optional) Agent trained\"".format(log_path))
Explanation: One last task before starting the training: let's upload the tensoboard logs, to get an overview of the performance of our model. We will upload our logs to tensorboard.dev and for that you need to 
run the print statement below and copy the output of the cell (which is a command) into a terminal, then execute the command from there. It will give you a link from which you need to copy/paste the authentication code, and once that is done, you will receive the 
url of your model evaluation, hosted on a public tensorboard.dev instance. As soon as you kicked off the training in the subsequent cell, you should see some graphs as in the picture below.
End of explanation
import numpy as np
feature = np.reshape(environment._observe()[0], (1,20))
feature.shape
## Inference
step = ts.TimeStep(
        tf.constant(
            ts.StepType.FIRST, dtype=tf.int32, shape=[1],
            name='step_type'),
        tf.constant(0.0, dtype=tf.float32, shape=[1], name='reward'),
        tf.constant(1.0, dtype=tf.float32, shape=[1], name='discount'),
        tf.constant(feature,
                    dtype=tf.float64, shape=[1, 20],
                    name='observation'))
agent.policy.action(step).action.numpy()
Explanation: <img src='./assets/example_tensorboard.png'>
7. Inferencing with trained model & Tensorboard Evaluation
Now that our model is trained, what if we want to determine which action to take given a new "context": for that we will iterate on our dataset to get the next item,
    make a timestep out of it by wrapping the results using ts.Timestep. It expects step_type, reward, discount, and observation as input: since we are performing prediction you can fill 
        in dummy values for the first 3: only the observation/context is relevant. Read about how it works here, and perform the task below
the movielens environment provides us a private observe_method which randomly samples upto 8 user context observations, and we select one of them, and reshape it to (1,20): the shape required for the model to consume.
End of explanation
<END_TASK>