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
extern crate cpython; use cpython::{ PyObject, ObjectProtocol, ToPyObject, NoArgs }; pub use cpython::{ Python, PyResult, PyDict }; macro_rules! call { ($obj:expr, $py:expr, $name:expr, $args:expr) => { $obj.call_method($py, $name, $args, None).unwrap() } } pub struct Env<'a> { obj: PyObject, py: Python<'a> } impl<'a> Env<'a> { pub fn make(py: Python<'a>, game_id: &str) -> Env<'a> { py.run("import sys\nsys.argv = ['']", None, None).unwrap(); Env { obj: py.import("gym").unwrap() .call(py, "make", (game_id, ), None).unwrap(), py } } pub fn reset(&self) -> Vec<f64> { let multi_list: PyObject = call!(self.obj, self.py, "reset", NoArgs {}); self.py.import("numpy").unwrap() .call(self.py, "ravel", (multi_list, ), None).unwrap() .extract(self.py).unwrap() } pub fn step<T: ToPyObject>(&self, action: T) -> (Vec<f64>, f64, bool, PyDict) { call!(self.obj, self.py, "step", (action,)).extract(self.py).unwrap() } pub fn render(&self) { call!(self.obj, self.py, "render", NoArgs); } } #[cfg(test)] mod tests { use ::Env; use cpython::Python; #[test] fn make_env() { let gil = Python::acquire_gil(); let env = Env::make(gil.python(), "LunarLander-v2"); env.reset(); for _ in 0..1000 { env.render(); if env.step(1).2 { env.reset(); } } } }