rlmesh-runtime 0.1.0-beta.1

Runtime driver for RLMesh model-environment evaluation sessions
Documentation 
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
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76

use std::collections::BTreeMap;

#[derive(Debug, Clone)]
pub(crate) struct EpisodeRecord {
    pub(crate) record_id: String,
    pub(crate) index: i64,
    pub(crate) env_index: i32,
    pub(crate) started_from_auto_reset: bool,
}

#[derive(Debug, Default)]
pub(crate) struct EpisodeRecordRegistry {
    next_index: i64,
    by_episode_id: BTreeMap<String, EpisodeRecord>,
}

impl EpisodeRecordRegistry {
    pub(crate) fn ensure_for_slots(
        &mut self,
        episode_ids: &[String],
        started_from_auto_reset: bool,
    ) -> (Vec<String>, Vec<(String, EpisodeRecord)>) {
        let mut record_ids = Vec::with_capacity(episode_ids.len());
        let mut started = Vec::new();
        for (env_index, episode_id) in episode_ids.iter().enumerate() {
            if episode_id.is_empty() {
                record_ids.push(String::new());
                continue;
            }
            if let Some(existing) = self.by_episode_id.get(episode_id) {
                record_ids.push(existing.record_id.clone());
                continue;
            }

            self.next_index += 1;
            let record = EpisodeRecord {
                record_id: format!("ep-{:06}", self.next_index),
                index: self.next_index,
                env_index: env_index as i32,
                started_from_auto_reset,
            };
            record_ids.push(record.record_id.clone());
            self.by_episode_id
                .insert(episode_id.clone(), record.clone());
            started.push((episode_id.clone(), record));
        }
        (record_ids, started)
    }

    pub(crate) fn record_for(&self, episode_id: &str) -> Option<&EpisodeRecord> {
        self.by_episode_id.get(episode_id)
    }
}

#[cfg(test)]
mod tests {
    use super::EpisodeRecordRegistry;

    #[test]
    fn episode_record_ids_are_global_across_vector_slots() {
        let mut registry = EpisodeRecordRegistry::default();
        let (initial_ids, started) =
            registry.ensure_for_slots(&["runtime-a".to_string(), "runtime-b".to_string()], false);
        assert_eq!(initial_ids, ["ep-000001", "ep-000002"]);
        assert_eq!(started.len(), 2);
        assert!(!started[0].1.started_from_auto_reset);

        let (next_ids, started) =
            registry.ensure_for_slots(&["runtime-a".to_string(), "runtime-c".to_string()], true);
        assert_eq!(next_ids, ["ep-000001", "ep-000003"]);
        assert_eq!(started.len(), 1);
        assert_eq!(started[0].0, "runtime-c");
        assert_eq!(started[0].1.env_index, 1);
        assert!(started[0].1.started_from_auto_reset);
    }
}