reinforcex 0.0.4

Deep Reinforcement Learning Framework
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56

use super::experience::Experience;
use crate::misc::bounded_vec_deque::BoundedVecDeque;
use crate::prob_distributions::BaseDistribution;
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use tch::Tensor;
use ulid::Ulid;

pub struct OnPolicyBuffer {
    experiences_by_episode: HashMap<Ulid, BoundedVecDeque<Arc<Experience>>>,
}

impl OnPolicyBuffer {
    pub fn new() -> Self {
        Self {
            experiences_by_episode: HashMap::new(),
        }
    }

    pub fn append(
        &mut self,
        agent_id: Ulid,
        episode_id: Ulid,
        state: Tensor,
        action: Option<Tensor>,
        action_distrib: Option<Box<dyn BaseDistribution>>,
        reward: f64,
        is_episode_terminal: bool,
    ) -> Arc<Experience> {
        let experience = Arc::new(Experience::new(
            agent_id,
            episode_id,
            state,
            action,
            action_distrib,
            reward,
            is_episode_terminal,
            Mutex::new(None),
            Mutex::new(None),
        ));

        self.experiences_by_episode
            .entry(episode_id)
            .or_insert_with(|| BoundedVecDeque::new(1e9 as usize))
            .push_back(experience.clone());

        experience
    }

    pub fn flush(&mut self) -> Vec<Vec<Arc<Experience>>> {
        self.experiences_by_episode
            .drain()
            .map(|(_k, v)| v.to_vec())
            .collect()
    }
}