use crate::env::Environment;
pub struct Exp<E: Environment> {
pub state: E::State,
pub action: E::Action,
pub next_state: Option<E::State>,
pub reward: f32,
}
impl<E: Environment> Clone for Exp<E> {
fn clone(&self) -> Self {
Self {
state: self.state.clone(),
action: self.action.clone(),
next_state: self.next_state.clone(),
reward: self.reward,
}
}
}
#[derive(Clone, Debug)]
pub struct ExpBatch<E: Environment> {
pub states: Vec<E::State>,
pub actions: Vec<E::Action>,
pub next_states: Vec<Option<E::State>>,
pub rewards: Vec<f32>,
}
impl<E: Environment> ExpBatch<E> {
pub fn from_iter(iter: impl IntoIterator<Item = Exp<E>>, batch_size: usize) -> Self {
let batch = Self {
states: Vec::with_capacity(batch_size),
actions: Vec::with_capacity(batch_size),
next_states: Vec::with_capacity(batch_size),
rewards: Vec::with_capacity(batch_size),
};
iter.into_iter().fold(batch, |mut b, e| {
b.states.push(e.state.clone());
b.actions.push(e.action.clone());
b.next_states.push(e.next_state.clone());
b.rewards.push(e.reward);
b
})
}
}
#[cfg(test)]
mod tests {
use crate::env::tests::MockEnv;
use super::*;
const BATCH_SIZE: usize = 2;
fn create_mock_exp_array() -> [Exp<MockEnv>; BATCH_SIZE] {
let exp1 = Exp {
state: 0,
action: 1,
next_state: Some(1),
reward: 1.0,
};
let exp2 = Exp {
state: 1,
action: 2,
next_state: None,
reward: 0.0,
};
[exp1, exp2]
}
#[test]
fn exp_batch_from_iter() {
let experiences = create_mock_exp_array();
let batch = ExpBatch::from_iter(experiences, BATCH_SIZE);
assert_eq!(batch.states, [0, 1], "States constructed correctly");
assert_eq!(batch.actions, [1, 2], "States constructed correctly");
assert_eq!(
batch.next_states,
[Some(1), None],
"States constructed correctly"
);
assert_eq!(batch.rewards, [1.0, 0.0], "States constructed correctly");
}
}