bevy_rl 0.9.0

Build environments for reinforcement learning with bevy
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
132
133

use std::sync::{Arc, Mutex};

use image;

use bevy::prelude::*;
use crossbeam_channel::*;

use crate::render;

#[derive(Resource)]
pub struct AIGymStateInner<
    A: 'static + Send + Sync + Clone + std::panic::RefUnwindSafe,
    B: 'static + Send + Sync + Clone + std::panic::RefUnwindSafe,
> {
    // Bevy image handle for the screen
    pub render_image_handles: Vec<Handle<Image>>,

    // Sync with engine thread.
    pub(crate) _step_tx: Sender<Vec<Option<String>>>,
    pub(crate) _step_rx: Receiver<Vec<Option<String>>>,

    pub(crate) _reset_tx: Sender<bool>,
    pub(crate) _reset_rx: Receiver<bool>,

    pub(crate) _step_result_tx: Sender<Vec<bool>>,
    pub(crate) _step_result_rx: Receiver<Vec<bool>>,

    pub(crate) _reset_result_tx: Sender<bool>,
    pub(crate) _reset_result_rx: Receiver<bool>,

    pub(crate) _environment_state: Option<B>,

    // State
    pub visual_observations: Vec<image::RgbaImage>,
    pub alt_observation: Option<image::RgbaImage>,
    pub rewards: Vec<f32>,
    pub actions: Vec<Option<A>>,
    pub terminations: Vec<bool>,
}

impl<
        A: 'static + Send + Sync + Clone + std::panic::RefUnwindSafe,
        B: 'static + Send + Sync + Clone + std::panic::RefUnwindSafe + serde::Serialize,
    > AIGymStateInner<A, B>
{
    pub fn new(settings: render::AIGymSettings) -> Self {
        let (step_tx, step_rx) = bounded(1);
        let (reset_tx, reset_rx) = bounded(1);
        let (result_tx, result_rx) = bounded(1);
        let (result_reset_tx, result_reset_rx) = bounded(1);
        Self {
            // Channels
            _step_tx: step_tx,
            _step_rx: step_rx,
            _step_result_tx: result_tx,
            _step_result_rx: result_rx,

            _reset_tx: reset_tx,
            _reset_rx: reset_rx,
            _reset_result_tx: result_reset_tx,
            _reset_result_rx: result_reset_rx,

            _environment_state: None,

            // Render Targets
            render_image_handles: Vec::new(),

            // State
            visual_observations: Vec::new(),
            alt_observation: None,
            rewards: vec![0.0; settings.num_agents as usize],
            actions: vec![None; settings.num_agents as usize],
            terminations: vec![false; settings.num_agents as usize],
        }
    }

    pub fn send_step_result(&self, results: Vec<bool>) {
        if self._step_result_tx.is_empty() {
            self._step_result_tx.send(results).unwrap();
        }
    }

    pub fn send_reset_result(&self, result: bool) {
        if self._reset_result_tx.is_empty() {
            self._reset_result_tx.send(result).unwrap();
        }
    }

    pub fn receive_action_strings(&self) -> Vec<Option<String>> {
        self._step_rx.recv().unwrap()
    }

    pub fn receive_reset_request(&self) {
        self._reset_rx.recv().unwrap();
    }

    pub fn is_next_action(&self) -> bool {
        !self._step_rx.is_empty()
    }

    pub fn is_reset_request(&self) -> bool {
        !self._reset_rx.is_empty()
    }

    pub fn set_reward(&mut self, agent_index: usize, score: f32) {
        self.rewards[agent_index] = score;
    }

    pub fn set_terminated(&mut self, agent_index: usize, result: bool) {
        self.terminations[agent_index] = result;
    }

    pub fn reset(&mut self) {
        for i in 0..self.terminations.len() {
            self.set_terminated(i, false);
            self.set_reward(i, 0.0);
        }

        self.send_reset_result(true);
    }

    pub fn set_env_state(&mut self, state: B) {
        self._environment_state = Some(state);
    }
}

#[derive(Resource, Clone)]
pub struct AIGymState<
    A: 'static + Send + Sync + Clone + std::panic::RefUnwindSafe,
    B: 'static + Send + Sync + Clone + std::panic::RefUnwindSafe,
> {
    pub guard: Arc<Mutex<AIGymStateInner<A, B>>>,
}