80 lines
2.0 KiB
Rust
80 lines
2.0 KiB
Rust
// Minimal test to see if we can run any RL tests at all
|
|
use crate::{Environment, Step};
|
|
use std::collections::HashMap;
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[derive(Clone, Debug)]
|
|
struct DummyEnv {
|
|
state: f32,
|
|
step_count: usize,
|
|
}
|
|
|
|
impl Environment for DummyEnv {
|
|
type State = f32;
|
|
type Action = f32;
|
|
type Reward = f32;
|
|
|
|
fn reset(&mut self) -> Self::State {
|
|
self.state = 0.0;
|
|
self.step_count = 0;
|
|
self.state
|
|
}
|
|
|
|
fn step(&mut self, action: &Self::Action) -> Step<Self::State, Self::Reward> {
|
|
self.state += action;
|
|
self.step_count += 1;
|
|
|
|
Step {
|
|
state: self.state,
|
|
reward: -self.state.abs(),
|
|
done: self.step_count >= 10,
|
|
info: HashMap::new(),
|
|
}
|
|
}
|
|
|
|
fn action_space(&self) -> (Self::Action, Self::Action) {
|
|
(-1.0, 1.0)
|
|
}
|
|
|
|
fn observation_space(&self) -> (Self::State, Self::State) {
|
|
(-10.0, 10.0)
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_simple_environment() {
|
|
let mut env = DummyEnv { state: 0.0, step_count: 0 };
|
|
|
|
let initial_state = env.reset();
|
|
assert_eq!(initial_state, 0.0);
|
|
|
|
let step_result = env.step(&0.5);
|
|
assert_eq!(step_result.state, 0.5);
|
|
assert_eq!(step_result.reward, -0.5);
|
|
assert!(!step_result.done);
|
|
|
|
let (action_low, action_high) = env.action_space();
|
|
assert_eq!(action_low, -1.0);
|
|
assert_eq!(action_high, 1.0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_environment_termination() {
|
|
let mut env = DummyEnv { state: 0.0, step_count: 0 };
|
|
env.reset();
|
|
|
|
// Step until termination
|
|
for i in 0..15 {
|
|
let step_result = env.step(&0.1);
|
|
if i >= 9 {
|
|
assert!(step_result.done);
|
|
} else {
|
|
assert!(!step_result.done);
|
|
}
|
|
}
|
|
}
|
|
}
|