optionally invert agent-env relation

src/agent.rs

@@ -0,0 +1,9 @@
+use crate::{env::Environment, memory::Exp};
+
+pub trait Agent<E>
+where
+    E: Environment,
+{
+    fn act(&self, state: &E::State, actions: &[E::Action]) -> E::Action;
+    fn learn(&mut self, exp: Exp<E>, next_actions: &[E::Action]);
+}

src/algo/tabular/action_occurrence.rs

@@ -1,6 +1,9 @@
 use std::collections::HashMap;
 
+use rand::Rng;
+
 use crate::{
+    agent::Agent,
     decay::{self, Decay},
     env::{DiscreteActionSpace, Environment},
     exploration::{Choice, EpsilonGreedy},
@@ -157,3 +160,39 @@ where
         self.episode += 1;
     }
 }
+impl<E, D> Agent<E> for ActionOccurrenceAgent<E, D>
+where
+    E: Environment + DiscreteActionSpace,
+    E::State: Hashable,
+    E::Action: Hashable,
+    D: Decay,
+{
+    fn act(&self, state: &E::State, actions: &[E::Action]) -> E::Action {
+        match self.exploration.choose(self.episode) {
+            Choice::Explore => actions[rand::thread_rng().gen_range(0..actions.len())],
+            Choice::Exploit => *actions
+                .iter()
+                .max_by(|&a, &b| {
+                    let a_value = self
+                        .table
+                        .get(&(*state, *a))
+                        .copied()
+                        .unwrap_or_default()
+                        .value;
+                    let b_value = self
+                        .table
+                        .get(&(*state, *b))
+                        .copied()
+                        .unwrap_or_default()
+                        .value;
+                    a_value.partial_cmp(&b_value).unwrap()
+                })
+                .expect("There is always at least one action available"),
+        }
+    }
+
+    fn learn(&mut self, exp: Exp<E>, _: &[E::Action]) {
+        self.learn(exp);
+        self.episode += 1;
+    }
+}

src/algo/tabular/q_table.rs

@@ -1,6 +1,9 @@
 use std::collections::HashMap;
 
+use rand::Rng;
+
 use crate::{
+    agent::Agent,
     assert_interval, decay,
     env::{DiscreteActionSpace, Environment},
     exploration::{Choice, EpsilonGreedy},
@@ -142,3 +145,31 @@ where
         self.episode += 1;
     }
 }
+
+impl<E> Agent<E> for QTableAgent<E>
+where
+    E: Environment + DiscreteActionSpace,
+    E::State: Hashable,
+    E::Action: Hashable,
+{
+    fn act(&self, state: &E::State, actions: &[E::Action]) -> <E as Environment>::Action {
+        match self.exploration.choose(self.episode) {
+            Choice::Explore => actions[rand::thread_rng().gen_range(0..actions.len())],
+            Choice::Exploit => {
+                *actions
+                    .iter()
+                    .max_by(|&a, &b| {
+                        let a_value = *self.q_table.get(&(*state, *a)).unwrap_or(&0.0);
+                        let b_value = *self.q_table.get(&(*state, *b)).unwrap_or(&0.0);
+                        a_value.partial_cmp(&b_value).unwrap()
+                    })
+                    .expect("There is always at least one action available") // Maybe make this more lenient by providing a default?
+            }
+        }
+    }
+
+    fn learn(&mut self, exp: Exp<E>, next_actions: &[E::Action]) {
+        self.learn(exp, next_actions);
+        self.episode += 1;
+    }
+}

src/algo/tabular/ucb.rs

@@ -1,6 +1,7 @@
 use std::collections::HashMap;
 
 use crate::{
+    agent::Agent,
     env::{DiscreteActionSpace, Environment},
     memory::Exp,
 };
@@ -153,3 +154,17 @@ where
         self.episode += 1;
     }
 }
+impl<E> Agent<E> for UCBAgent<E>
+where
+    E: Environment + DiscreteActionSpace,
+    E::State: Hashable,
+    E::Action: Hashable + From<usize>,
+{
+    fn act(&self, state: &E::State, actions: &[E::Action]) -> E::Action {
+        self.act(*state, actions)
+    }
+
+    fn learn(&mut self, exp: Exp<E>, _: &[E::Action]) {
+        self.learn(exp);
+    }
+}

src/env.rs

@@ -2,9 +2,10 @@ use std::{
     collections::BTreeMap,
     fmt::Debug,
     ops::{Deref, DerefMut},
+    os::unix::net,
 };
 
-use crate::util::summary_from_keys;
+use crate::{agent::Agent, memory::Exp, util::summary_from_keys};
 
 /// Represents a Markov decision process, defining the dynamics of an environment
 /// in which an agent can operate.
@@ -52,6 +53,31 @@ pub trait DiscreteActionSpace: Environment {
     ///
     /// The returned slice should never be empty, instead specify an action that represents doing nothing if necessary.
     fn actions(&self) -> Vec<Self::Action>;
+
+    fn run(&mut self, agent: &mut impl Agent<Self>)
+    where
+        Self: Sized,
+        Self::Action: Copy,
+        Self::State: Copy,
+    {
+        let mut next_state = Some(self.reset());
+        let mut actions = self.actions();
+        while let Some(state) = next_state {
+            let action = agent.act(&state, &actions);
+            let (next, reward) = self.step(action);
+            next_state = next;
+            actions = self.actions();
+            agent.learn(
+                Exp {
+                    state,
+                    action,
+                    reward,
+                    next_state,
+                },
+                &actions,
+            );
+        }
+    }
 }
 
 /// An [Environment] with a discrete state space

src/lib.rs

@@ -10,6 +10,8 @@ pub mod ds;
 /// Environment
 pub mod env;
 
+pub mod agent;
+
 /// Exploration policies
 pub mod exploration;