entity_gym_rs/low_level/
env.rs

1use rustc_hash::FxHashMap;
2
3pub type EntityId = u64;
4pub type ActionType = String;
5pub type EntityType = String;
6
7// Could have a `SingleAgentEnv` that implement `Environment`
8// Could maybe use const generic for `agents`, return fixed size slices?
9pub trait Environment {
10    fn obs_space(&self) -> ObsSpace;
11    fn action_space(&self) -> Vec<(ActionType, ActionSpace)>;
12    fn agents(&self) -> usize;
13
14    #[allow(clippy::vec_box)]
15    fn reset(&mut self) -> Vec<Box<Observation>>;
16    #[allow(clippy::vec_box)]
17    fn act(&mut self, action: &[Vec<Option<Action>>]) -> Vec<Box<Observation>>;
18    fn close(&mut self) {}
19}
20
21#[derive(Debug, Clone)]
22pub enum ActionSpace {
23    Categorical { choices: Vec<String> },
24    SelectEntity,
25}
26
27#[derive(Debug, Clone)]
28pub enum ActionMask {
29    DenseCategorical {
30        actors: Vec<EntityId>,
31        mask: Option<Vec<bool>>,
32    },
33    SelectEntity {
34        actors: Vec<EntityId>,
35        actees: Vec<EntityId>,
36    },
37}
38
39#[derive(Debug, Clone)]
40pub struct ObsSpace {
41    pub entities: Vec<(EntityType, Entity)>,
42}
43
44#[derive(Debug, Clone)]
45pub enum Action {
46    Categorical {
47        actors: Vec<EntityId>,
48        action: Vec<usize>,
49    },
50    SelectEntity {
51        actors: Vec<EntityId>,
52        actees: Vec<EntityId>,
53    },
54}
55
56#[derive(Debug, Clone)]
57pub struct Entity {
58    pub features: Vec<String>,
59}
60
61#[derive(Debug, Clone)]
62pub struct CompactFeatures {
63    pub counts: Vec<usize>,
64    pub data: Vec<f32>,
65}
66
67#[derive(Debug, Clone)]
68pub struct Observation {
69    pub features: CompactFeatures,
70    // Maps each player to (optional) list of IDs for all entities
71    pub ids: Vec<Option<Vec<EntityId>>>,
72    pub actions: Vec<Option<ActionMask>>,
73
74    pub done: bool,
75    pub reward: f32,
76    pub metrics: FxHashMap<String, f32>,
77}
entity_gym_rs/low_level/env.rs

entity_gym_rs/low_level/
env.rs
