bevy_rl 0.0.3

Using bevy to build gym environments to train reinforcement learning agents
Documentation 
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
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131

use crossbeam_channel::*;

use gotham::helpers::http::response::create_response;
use gotham::middleware::state::StateMiddleware;
use gotham::pipeline::{single_middleware, single_pipeline};
use gotham::router::builder::*;
use gotham::router::Router;
use gotham::state::StateData;
use gotham::state::{FromState, State};
use hyper::{body, Body, Response, StatusCode};

use futures::executor;
use image;
use std::io::Cursor;
use std::sync::{Arc, Mutex};

use crate::state;

// ---------------
// AI Gym REST API
// ---------------

#[derive(Clone, StateData)]
pub(crate) struct GothamState<T: 'static + Send + Sync + Clone + std::panic::RefUnwindSafe> {
    pub(crate) inner: Arc<Mutex<state::AIGymState<T>>>,
}

pub(crate) fn router<T: 'static + Send + Sync + Clone + std::panic::RefUnwindSafe>(
    state: GothamState<T>,
) -> Router {
    let middleware = StateMiddleware::new(state);
    let pipeline = single_middleware(middleware);

    let (chain, pipelines) = single_pipeline(pipeline);

    // build a router with the chain & pipeline
    build_router(chain, pipelines, |route| {
        route.get("/screen.png").to(screen::<T>);
        route.post("/step").to(step::<T>);
        route.post("/reset").to(reset::<T>);
    })
}

fn screen<T: 'static + Send + Sync + Clone + std::panic::RefUnwindSafe>(
    state: State,
) -> (State, Response<Body>) {
    let mut bytes: Vec<u8> = Vec::new();
    {
        let state_: &GothamState<T> = GothamState::borrow_from(&state);
        let state__ = state_.inner.lock().unwrap();
        let image = state__.screen.clone().unwrap();

        image
            .write_to(&mut Cursor::new(&mut bytes), image::ImageOutputFormat::Png)
            .unwrap();
    }
    let response = create_response::<Vec<u8>>(&state, StatusCode::OK, mime::TEXT_PLAIN, bytes);

    return (state, response);
}

fn step<T: 'static + Send + Sync + Clone + std::panic::RefUnwindSafe>(
    mut state: State,
) -> (State, String) {
    let body_ = Body::take_from(&mut state);
    let valid_body = executor::block_on(body::to_bytes(body_)).unwrap();
    let action = String::from_utf8(valid_body.to_vec()).unwrap();

    let state_: &GothamState<T> = GothamState::borrow_from(&state);
    let step_tx: Sender<String>;
    let result_rx: Receiver<bool>;

    let is_terminated: bool;
    {
        let ai_gym_state = state_.inner.lock().unwrap();
        is_terminated = ai_gym_state.is_terminated;
        step_tx = ai_gym_state.__step_channel_tx.clone();
        result_rx = ai_gym_state.__result_channel_rx.clone();
    }

    if is_terminated {
        return (
            state,
            format!("{{\"reward\": {}, \"is_terminated\": {}}}", 0, true),
        );
    }

    step_tx.send(action).unwrap();
    result_rx.recv().unwrap();

    let mut reward = 0.0;
    let is_terminated;
    {
        let ai_gym_state = state_.inner.lock().unwrap();

        if ai_gym_state.rewards.len() > 0 {
            reward = ai_gym_state.rewards[ai_gym_state.rewards.len() - 1];
        }
        if ai_gym_state.rewards.len() > 1 {
            reward -= ai_gym_state.rewards[ai_gym_state.rewards.len() - 2];
        }

        is_terminated = ai_gym_state.is_terminated.clone();
    }

    return (
        state,
        format!(
            "{{\"reward\": {}, \"is_terminated\": {}}}",
            reward, is_terminated
        ),
    );
}

fn reset<T: 'static + Send + Sync + Clone + std::panic::RefUnwindSafe>(
    state: State,
) -> (State, String) {
    let reset_channel_tx: Sender<bool>;
    let reset_result_channel_rx: Receiver<bool>;
    {
        let state_: &GothamState<T> = GothamState::borrow_from(&state);
        let ai_gym_state = state_.inner.lock().unwrap();
        reset_channel_tx = ai_gym_state.__reset_channel_tx.clone();
        reset_result_channel_rx = ai_gym_state.__result_reset_channel_rx.clone();
    }

    reset_channel_tx.send(true).unwrap();
    reset_result_channel_rx.recv().unwrap();

    return (state, "ok".to_string());
}