entity_gym_rs/agent/
mod.rs

1mod action;
2mod featurizable;
3mod obs;
4mod random;
5mod rogue_net;
6#[cfg(feature = "bevy")]
7mod rogue_net_asset;
8#[cfg(feature = "python")]
9mod training;
10
11use std::io::Read;
12use std::path::Path;
13use std::sync::atomic::{AtomicUsize, Ordering};
14use std::sync::Arc;
15
16pub use self::rogue_net::RogueNetAgent;
17pub use action::Action;
18use crossbeam_channel::Receiver;
19pub use entity_gym_derive::*;
20pub use featurizable::Featurizable;
21pub use obs::Obs;
22pub use random::RandomAgent;
23#[cfg(feature = "bevy")]
24pub use rogue_net_asset::{RogueNetAsset, RogueNetAssetLoader};
25#[cfg(feature = "python")]
26pub use training::{TrainAgent, TrainAgentEnv, TrainEnvBuilder};
27
28/// Agents are given observations and return actions.
29///
30/// You don't generally need to implement the [`Agent`] trait yourself.
31/// There are three main ways of obtaining an [`Agent`]:
32/// 1. [`random()`] creates an agent that chooses actions uniformly at random.
33/// 2. [`load`] and [`load_archive`] loads a trained neural network agent from an [enn-trainer](https://github.com/entity-neural-network/enn-trainer) checkpoint directory or an archive of a checkpoint directory.
34/// 3. [`TrainEnvBuilder`] can be used to obtain a [`TrainAgent`]/[`TrainAgentEnv`] pair which can be used to train a neural network agent.
35///
36/// Every [`Agent`] also implements the [`AgentOps`] trait which provides more ergonomic typed versions of the [`Agent::act_dyn`] and [`Agent::act_async_dyn`] methods.
37pub trait Agent {
38    /// Returns an action for the given observation.
39    fn act_dyn(&mut self, action: &str, num_actions: u64, obs: &Obs) -> Option<Vec<u64>>;
40
41    /// Returns receiver that can be blocked on to receive an action for the given observation.
42    #[must_use]
43    fn act_async_dyn(&mut self, action: &str, num_actions: u64, obs: &Obs) -> ActionReceiver<u64>;
44
45    /// Indicates that the agent has reached the end of the training episode.
46    fn game_over(&mut self, obs: &Obs);
47}
48
49/// Augments the [`Agent`] trait with more ergonomic typed versions of the [`Agent::act_dyn`] and [`Agent::act_async_dyn`] methods.
50pub trait AgentOps {
51    /// Returns an action for the given observation.
52    fn act<'a, A: Action<'a>>(&mut self, obs: &Obs) -> Option<Vec<A>>;
53
54    /// Returns receiver that can be blocked on to receive an action for the given observation.
55    #[must_use]
56    fn act_async<'a, A: Action<'a>>(&mut self, obs: &Obs) -> ActionReceiver<A>;
57}
58
59impl<T: Agent> AgentOps for T {
60    fn act<'a, A: Action<'a>>(&mut self, obs: &Obs) -> Option<Vec<A>> {
61        self.act_dyn(A::name(), A::num_actions(), obs)
62            .map(|x| x.into_iter().map(A::from_u64).collect())
63    }
64
65    #[must_use]
66    fn act_async<'a, A: Action<'a>>(&mut self, obs: &Obs) -> ActionReceiver<A> {
67        let receiver = self.act_async_dyn(A::name(), A::num_actions(), obs);
68        unsafe { std::mem::transmute::<ActionReceiver<u64>, ActionReceiver<A>>(receiver) }
69    }
70}
71
72impl AgentOps for dyn Agent {
73    fn act<'a, A: Action<'a>>(&mut self, obs: &Obs) -> Option<Vec<A>> {
74        self.act_dyn(A::name(), A::num_actions(), obs)
75            .map(|x| x.into_iter().map(A::from_u64).collect())
76    }
77
78    #[must_use]
79    fn act_async<'a, A: Action<'a>>(&mut self, obs: &Obs) -> ActionReceiver<A> {
80        let receiver = self.act_async_dyn(A::name(), A::num_actions(), obs);
81        unsafe { std::mem::transmute::<ActionReceiver<u64>, ActionReceiver<A>>(receiver) }
82    }
83}
84
85/// A channel for receiving an agent action returned by [`AgentOps::act_async`] or [`Agent::act_async_dyn`].
86pub struct ActionReceiver<A> {
87    inner: InnerActionReceiver<A>,
88}
89
90enum InnerActionReceiver<A> {
91    // Variant is only constructed when cfg(feature = "python").
92    #[allow(dead_code)]
93    Receiver {
94        receiver: Receiver<Vec<u64>>,
95        observations_remaining: Arc<AtomicUsize>,
96        agent_count: usize,
97        phantom: std::marker::PhantomData<A>,
98    },
99    Value(Vec<u64>),
100}
101
102impl<A> ActionReceiver<A> {
103    /// Blocks on the receiver until an action is received.
104    pub fn rcv_raw(self) -> Option<Vec<u64>> {
105        match self.inner {
106            InnerActionReceiver::Receiver {
107                receiver,
108                observations_remaining,
109                agent_count,
110                ..
111            } => {
112                let remaining = observations_remaining.load(Ordering::SeqCst);
113                if remaining != 0 && (remaining != agent_count || agent_count == 1) {
114                    panic!("TrainAgent::act called before all agents have received observations. This is not allowed. If you have multiple agents, call the `act_async` on every agent before awaiting any actions.");
115                }
116                let act = receiver.recv();
117                observations_remaining.store(agent_count, Ordering::SeqCst);
118                act.ok()
119            }
120            InnerActionReceiver::Value(value) => Some(value),
121        }
122    }
123
124    /// Blocks on the receiver until an action is received.
125    pub fn rcv<'a>(self) -> Option<Vec<A>>
126    where
127        A: Action<'a>,
128    {
129        self.rcv_raw()
130            .map(|x| x.into_iter().map(A::from_u64).collect())
131    }
132
133    /// Creates a new [`ActionReceiver`] which will return the given value.
134    pub(crate) fn value(val: Vec<u64>) -> ActionReceiver<A> {
135        ActionReceiver {
136            inner: InnerActionReceiver::Value(val),
137        }
138    }
139}
140
141/// Returns a boxed [`RandomAgent`].
142pub fn random() -> Box<dyn Agent> {
143    Box::new(RandomAgent::default())
144}
145
146/// Returns a boxed [`RandomAgent`] with the given seed.
147pub fn random_seeded(seed: u64) -> Box<dyn Agent> {
148    Box::new(RandomAgent::from_seed(seed))
149}
150
151/// Loads an agent from a checkpoint directory.
152pub fn load<P: AsRef<Path>>(path: P) -> Box<dyn Agent> {
153    Box::new(RogueNetAgent::load(path).unwrap())
154}
155
156/// Loads an agent from an archive of a checkpoint directory.
157pub fn load_archive<R: Read>(reader: R) -> Result<Box<dyn Agent>, std::io::Error> {
158    Ok(Box::new(RogueNetAgent::load_archive(reader)?))
159}
entity_gym_rs/agent/mod.rs

entity_gym_rs/agent/
mod.rs
