burn_train/learner/rl/env_runner/
base.rs

1use burn_core::Tensor;
2use burn_core::data::dataloader::Progress;
3use burn_core::tensor::Device;
4use burn_rl::Policy;
5use burn_rl::ToObservation;
6use burn_rl::Transition;
7use burn_rl::{Environment, EnvironmentInit};
8
9use crate::RLEvent;
10use crate::{
11    AgentEvaluationEvent, EpisodeSummary, EvaluationItem, EventProcessorTraining,
12    RLEventProcessorType,
13};
14use crate::{Interrupter, RLComponentsTypes};
15
16/// A trajectory, i.e. a list of ordered [TimeStep](TimeStep).
17#[derive(Clone, new)]
18pub struct Trajectory<S, A, C> {
19    /// A list of ordered [TimeStep](TimeStep)s.
20    pub timesteps: Vec<TimeStep<S, A, C>>,
21}
22
23/// A timestep debscribing an iteration of the state/decision process.
24#[derive(Clone)]
25pub struct TimeStep<S, A, C> {
26    /// The environment id.
27    pub env_id: usize,
28    /// The [burn_rl::Transition](burn_rl::Transition).
29    pub transition: Transition<S, A>,
30    /// True if the environment reaches a terminal state.
31    pub done: bool,
32    /// The running length of the current episode.
33    pub ep_len: usize,
34    /// The running cumulative reward.
35    pub cum_reward: f64,
36    /// The action's context for this timestep.
37    pub action_context: C,
38}
39
40pub(crate) type RLTimeStep<RLC> = TimeStep<
41    <RLC as RLComponentsTypes>::State,
42    <RLC as RLComponentsTypes>::Action,
43    <RLC as RLComponentsTypes>::ActionContext,
44>;
45
46pub(crate) type RLTrajectory<RLC> = Trajectory<
47    <RLC as RLComponentsTypes>::State,
48    <RLC as RLComponentsTypes>::Action,
49    <RLC as RLComponentsTypes>::ActionContext,
50>;
51
52/// Trait for a structure that implements an agent/environement interface.
53pub trait AgentEnvLoop<RLC: RLComponentsTypes> {
54    /// Run a certain number of timesteps.
55    ///
56    /// # Arguments
57    ///
58    /// * `num_steps` - The number of time_steps to run.
59    /// * `processor` - An [crate::EventProcessorTraining](crate::EventProcessorTraining).
60    /// * `interrupter` - An [crate::Interrupter](crate::Interrupter).
61    /// * `num_steps` - The number of time_steps to run.
62    /// * `progress` - A mutable reference to the learning progress.
63    ///
64    /// # Returns
65    ///
66    /// A list of ordered timesteps.
67    fn run_steps(
68        &mut self,
69        num_steps: usize,
70        processor: &mut RLEventProcessorType<RLC>,
71        interrupter: &Interrupter,
72        progress: &mut Progress,
73    ) -> Vec<RLTimeStep<RLC>>;
74    /// Run a certain number of episodes.
75    ///
76    /// # Arguments
77    ///
78    /// * `num_episodes` - The number of episodes to run.
79    /// * `processor` - An [crate::EventProcessorTraining](crate::EventProcessorTraining).
80    /// * `interrupter` - An [crate::Interrupter](crate::Interrupter).
81    /// * `progress` - A mutable reference to the learning progress.
82    ///
83    /// # Returns
84    ///
85    /// A list of ordered timesteps.
86    fn run_episodes(
87        &mut self,
88        num_episodes: usize,
89        processor: &mut RLEventProcessorType<RLC>,
90        interrupter: &Interrupter,
91        progress: &mut Progress,
92    ) -> Vec<RLTrajectory<RLC>>;
93    /// Update the runner's agent.
94    fn update_policy(&mut self, update: RLC::PolicyState);
95    /// Get the state of the runner's agent.
96    fn policy(&self) -> RLC::PolicyState;
97    /// Returns the device on which the runner's agent runs.
98    fn device(&self) -> Device;
99}
100
101/// A simple, synchronized agent/environement interface.
102pub struct AgentEnvBaseLoop<RLC: RLComponentsTypes> {
103    env: RLC::Env,
104    eval: bool,
105    agent: RLC::Policy,
106    deterministic: bool,
107    current_reward: f64,
108    run_num: usize,
109    step_num: usize,
110    device: Device,
111}
112
113impl<RLC: RLComponentsTypes> AgentEnvBaseLoop<RLC> {
114    /// Create a new base runner.
115    pub fn new(
116        env_init: RLC::EnvInit,
117        agent: RLC::Policy,
118        eval: bool,
119        deterministic: bool,
120        device: &Device,
121    ) -> Self {
122        let agent = agent.to_device(device);
123        let mut env = env_init.init();
124        env.reset();
125
126        Self {
127            env,
128            eval,
129            agent: agent.clone(),
130            deterministic,
131            current_reward: 0.0,
132            run_num: 0,
133            step_num: 0,
134            device: device.clone(),
135        }
136    }
137}
138
139impl<RLC> AgentEnvLoop<RLC> for AgentEnvBaseLoop<RLC>
140where
141    RLC: RLComponentsTypes,
142{
143    fn run_steps(
144        &mut self,
145        num_steps: usize,
146        processor: &mut RLEventProcessorType<RLC>,
147        interrupter: &Interrupter,
148        progress: &mut Progress,
149    ) -> Vec<RLTimeStep<RLC>> {
150        let mut items = vec![];
151        let device = Default::default();
152        for _ in 0..num_steps {
153            let state = self.env.state();
154            let (action, context) = self.agent.action(
155                state.clone().to_observation(&self.device),
156                self.deterministic,
157            );
158
159            let step_result = self.env.step(RLC::Action::from(action.clone()));
160
161            self.current_reward += step_result.reward;
162            self.step_num += 1;
163
164            let transition = Transition::new(
165                state.clone(),
166                step_result.next_state,
167                RLC::Action::from(action),
168                Tensor::from_data([step_result.reward], &device),
169                Tensor::from_data(
170                    [(step_result.done || step_result.truncated) as i32 as f64],
171                    &device,
172                ),
173            );
174            items.push(TimeStep {
175                env_id: 0,
176                transition,
177                done: step_result.done,
178                ep_len: self.step_num,
179                cum_reward: self.current_reward,
180                action_context: context[0].clone(),
181            });
182
183            if !self.eval {
184                progress.items_processed += 1;
185                interrupter.fail_on_error(processor.process_train(RLEvent::EnvStep(
186                    EvaluationItem::new(context[0].clone(), progress.clone(), None),
187                )));
188
189                if step_result.done {
190                    interrupter.fail_on_error(processor.process_train(RLEvent::EpisodeEnd(
191                        EvaluationItem::new(
192                            EpisodeSummary {
193                                episode_length: self.step_num,
194                                cum_reward: self.current_reward,
195                            },
196                            progress.clone(),
197                            None,
198                        ),
199                    )));
200                }
201            }
202
203            if interrupter.should_stop() {
204                break;
205            }
206
207            if step_result.done || step_result.truncated {
208                self.env.reset();
209                self.current_reward = 0.;
210                self.step_num = 0;
211                self.run_num += 1;
212            }
213        }
214        items
215    }
216
217    fn update_policy(&mut self, update: RLC::PolicyState) {
218        self.agent.update(update);
219    }
220
221    fn run_episodes(
222        &mut self,
223        num_episodes: usize,
224        processor: &mut RLEventProcessorType<RLC>,
225        interrupter: &Interrupter,
226        progress: &mut Progress,
227    ) -> Vec<RLTrajectory<RLC>> {
228        self.env.reset();
229
230        let mut items = vec![];
231        for ep in 0..num_episodes {
232            let mut steps = vec![];
233            loop {
234                let step = self.run_steps(1, processor, interrupter, progress)[0].clone();
235                steps.push(step.clone());
236
237                if self.eval {
238                    interrupter.fail_on_error(processor.process_valid(
239                        AgentEvaluationEvent::EnvStep(EvaluationItem::new(
240                            step.action_context.clone(),
241                            Progress::new(
242                                steps.len() + 1,
243                                steps.len() + 1,
244                                Some("steps".to_string()),
245                            ),
246                            None,
247                        )),
248                    ));
249
250                    if step.done {
251                        interrupter.fail_on_error(processor.process_valid(
252                            AgentEvaluationEvent::EpisodeEnd(EvaluationItem::new(
253                                EpisodeSummary {
254                                    episode_length: step.ep_len,
255                                    cum_reward: step.cum_reward,
256                                },
257                                Progress::new(ep + 1, num_episodes, Some("episodes".to_string())),
258                                None,
259                            )),
260                        ));
261                    }
262                }
263
264                if interrupter.should_stop() || step.done {
265                    break;
266                }
267            }
268            items.push(Trajectory::new(steps));
269
270            if interrupter.should_stop() {
271                break;
272            }
273        }
274        items
275    }
276
277    fn policy(&self) -> RLC::PolicyState {
278        self.agent.state()
279    }
280
281    fn device(&self) -> Device {
282        self.device.clone()
283    }
284}
285
286#[cfg(test)]
287#[allow(clippy::needless_range_loop)]
288mod tests {
289    use crate::AsyncProcessorTraining;
290
291    use crate::learner::tests::{
292        MockEnvInit, MockPolicy, MockPolicyState, MockProcessor, MockRLComponents,
293    };
294
295    use super::*;
296
297    fn setup(state: usize, eval: bool, deterministic: bool) -> AgentEnvBaseLoop<MockRLComponents> {
298        let env_init = MockEnvInit;
299        let agent = MockPolicy(state);
300        AgentEnvBaseLoop::<MockRLComponents>::new(
301            env_init,
302            agent,
303            eval,
304            deterministic,
305            &Default::default(),
306        )
307    }
308
309    #[test]
310    fn test_policy_returns_agent_state() {
311        let runner = setup(1000, false, false);
312        let policy_state = runner.policy();
313        assert_eq!(policy_state.0, 1000);
314    }
315
316    #[test]
317    fn test_update_policy() {
318        let mut runner = setup(0, false, false);
319
320        runner.update_policy(MockPolicyState(1));
321        assert_eq!(runner.policy().0, 1);
322    }
323
324    #[test]
325    fn run_steps_returns_requested_number() {
326        let mut runner = setup(0, false, false);
327        let mut processor = AsyncProcessorTraining::new(MockProcessor);
328        let interrupter = Interrupter::new();
329        let mut progress = Progress {
330            items_processed: 0,
331            items_total: 1,
332            unit: Some("steps".to_string()),
333        };
334
335        let steps = runner.run_steps(1, &mut processor, &interrupter, &mut progress);
336        assert_eq!(steps.len(), 1);
337        let steps = runner.run_steps(8, &mut processor, &interrupter, &mut progress);
338        assert_eq!(steps.len(), 8);
339    }
340
341    #[test]
342    fn run_episodes_returns_requested_number() {
343        let mut runner = setup(0, false, false);
344        let mut processor = AsyncProcessorTraining::new(MockProcessor);
345        let interrupter = Interrupter::new();
346        let mut progress = Progress {
347            items_processed: 0,
348            items_total: 1,
349            unit: Some("steps".to_string()),
350        };
351
352        let trajectories = runner.run_episodes(1, &mut processor, &interrupter, &mut progress);
353        assert_eq!(trajectories.len(), 1);
354        assert_ne!(trajectories[0].timesteps.len(), 0);
355        let trajectories = runner.run_episodes(8, &mut processor, &interrupter, &mut progress);
356        assert_eq!(trajectories.len(), 8);
357        for i in 0..8 {
358            assert_ne!(trajectories[i].timesteps.len(), 0);
359        }
360    }
361}
burn_train/learner/rl/env_runner/base.rs

burn_train/learner/rl/env_runner/
base.rs
