entity_gym_rs/agent/
rogue_net.rs

1use std::fs::File;
2
3use ndarray::Array2;
4use rogue_net::{FwdArgs, RogueNet};
5
6use super::obs::EntityFeatures;
7use super::{ActionReceiver, Agent};
8use super::{Featurizable, Obs};
9
10/// Agent that implements the [RogueNet entity neural network](https://github.com/entity-neural-network/rogue-net).
11/// Can be loaded from checkpoints produced by [enn-trainer](https://github.com/entity-neural-network/enn-trainer).
12#[derive(Clone)]
13pub struct RogueNetAgent {
14    pub(crate) net: RogueNet,
15}
16
17impl RogueNetAgent {
18    /// Loads a neural network agent from an [enn-trainer](https://github.com/entity-neural-network/enn-trainer) checkpoint directory.
19    pub fn load<P: AsRef<std::path::Path>>(path: P) -> Result<Self, std::io::Error> {
20        let path = path.as_ref();
21        match path.extension() {
22            Some(ext) if ext == "roguenet" => Ok(RogueNetAgent {
23                net: RogueNet::load_archive(File::open(path)?)?,
24            }),
25            _ => Ok(RogueNetAgent {
26                net: RogueNet::load(path),
27            }),
28        }
29    }
30
31    /// Loads a neural network agent from an archive of a checkpoint directory.
32    ///
33    /// You can use the rogue-net cli to create a tar archive of a checkpoint directory:
34    /// ```console
35    /// $ cargo install rogue-net
36    /// $ rogue-net archive --path path/to/checkpoint/dir
37    /// ```
38    pub fn load_archive<R: std::io::Read>(reader: R) -> Result<Self, std::io::Error> {
39        let net = RogueNet::load_archive(reader)?;
40        Ok(RogueNetAgent { net })
41    }
42
43    /// Adapts the network to a changed observation space.
44    ///
45    /// If you trained a network with a different observation space, you can adapt it to the new observation space.
46    /// For this to work, the new set of features of the observation space must be a superset of the old set of features.
47    pub fn with_feature_adaptor<E: Featurizable>(mut self) -> Self {
48        self.net = self.net.with_obs_filter(
49            [(E::name().to_string(), E::feature_names())]
50                .iter()
51                .cloned()
52                .collect(),
53        );
54        self
55    }
56}
57
58impl Agent for RogueNetAgent {
59    fn act_dyn(&mut self, _action: &str, _num_actions: u64, obs: &Obs) -> Option<Vec<u64>> {
60        let features = obs
61            .entities
62            .iter()
63            .map(
64                |(
65                    name,
66                    EntityFeatures {
67                        features,
68                        num_entities,
69                        num_features,
70                        ..
71                    },
72                )| {
73                    (
74                        name.to_string(),
75                        Array2::from_shape_vec((*num_entities, *num_features), features.clone())
76                            .unwrap(),
77                    )
78                },
79            )
80            .collect();
81        let actors = obs
82            .entities
83            .iter()
84            .filter_map(|(name, e)| {
85                if e.is_actor {
86                    Some(name.to_string())
87                } else {
88                    None
89                }
90            })
91            .collect();
92        let (_probs, acts) = self.net.forward(FwdArgs { features, actors });
93        Some(acts)
94    }
95
96    fn act_async_dyn(&mut self, action: &str, num_actions: u64, obs: &Obs) -> ActionReceiver<u64> {
97        ActionReceiver::value(self.act_dyn(action, num_actions, obs).unwrap())
98    }
99
100    fn game_over(&mut self, _: &Obs) {}
101}
entity_gym_rs/agent/rogue_net.rs

entity_gym_rs/agent/
rogue_net.rs
