use burn::tensor::{backend::Backend, Tensor};
pub struct UCB<const A: usize> {
c: f32,
counter: [f32; A],
}
impl<const A: usize> UCB<A> {
pub fn new(c: f32) -> Self {
Self {
c,
counter: [1.0; A],
}
}
pub fn choose(&mut self, t: f32, q_values: &[f32; A]) -> usize {
let k = self.c * t.log10().sqrt();
let choice = q_values
.iter()
.enumerate()
.map(|(i, x)| (i, x + k * self.counter[i].powf(-0.5)))
.max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap())
.map(|(i, _)| i)
.expect("`q_values` is not empty");
self.counter[choice] += 1.0;
choice
}
#[allow(unused)]
pub fn choose_from_tensor<B: Backend>(&self, t: f32, tensor: Tensor<B, 1>) -> usize {
todo!()
}
}