old_contents: from six.moves import range
  import matplotlib.pyplot as plt
  
  from striatum.storage import (
      MemoryHistoryStorage,
      MemoryModelStorage,
      MemoryActionStorage,
      Action,
  )
  from striatum.bandit import UCB1
  from striatum import simulation
  
  
  def main():
      context_dimension = 5
      action_storage = MemoryActionStorage()
      action_storage.add([Action(i) for i in range(5)])
  
      # Regret Analysis
      n_rounds = 10000
      context, desired_actions = simulation.simulate_data(
          n_rounds, context_dimension, action_storage, random_state=1)
      policy = UCB1(MemoryHistoryStorage(), MemoryModelStorage(),
                    action_storage)
  
      for t in range(n_rounds):
          history_id, action = policy.get_action(context[t], 1)
          action_id = action[0]['action'].id
          if desired_actions[t] != action_id:
              policy.reward(history_id, {action_id: 0})
          else:
              policy.reward(history_id, {action_id: 1})
  
      policy.plot_avg_regret()
      plt.show()
  
  
  if __name__ == '__main__':
      main()
  
new_contents: from six.moves import range
  import matplotlib.pyplot as plt
  
  from striatum.storage import (
      MemoryHistoryStorage,
      MemoryModelStorage,
      MemoryActionStorage,
      Action,
  )
  from striatum.bandit import UCB1
  from striatum import simulation
  
  
  def main():
      context_dimension = 5
      action_storage = MemoryActionStorage()
      action_storage.add([Action(i) for i in range(5)])
  
      # Regret Analysis
      n_rounds = 10000
      context, desired_actions = simulation.simulate_data(
          n_rounds, context_dimension, action_storage, random_state=1)
      policy = UCB1(MemoryHistoryStorage(), MemoryModelStorage(),
                    action_storage)
  
      for t in range(n_rounds):
          history_id, recommendation = policy.get_action(context[t])
          action_id = recommendation.action.id
          if desired_actions[t] != action_id:
              policy.reward(history_id, {action_id: 0})
          else:
              policy.reward(history_id, {action_id: 1})
  
      policy.plot_avg_regret()
      plt.show()
  
  
  if __name__ == '__main__':
      main()
  
subject: fix simulation_ucb1