thrust_rl/env/games/
continuous_lqr.rs

1//! Minimal continuous-action environment for trait-wiring proof.
2//!
3//! `ContinuousLqr` is a one-dimensional linear quadratic regulator.
4//! It exists to demonstrate that the [`crate::env::Environment`] trait
5//! accepts a non-`i64` action (see issue #61), not to be a useful
6//! training target.
7//!
8//! State:  `[position, velocity]`.
9//! Action: `Vec<f32>` of length 1 (a scalar force in `[-1.0, 1.0]`).
10//! Reward: `-position^2 - 0.01 * action^2` (penalise being away from
11//! the origin and using control effort).
12//!
13//! Episode terminates after `max_steps` steps (default 50).
14//!
15//! Continuous-action training algorithms (SAC, TD3, DDPG, PPO with
16//! Gaussian heads) are explicitly *out of scope* for this env. Its
17//! only job is to compile against the trait.
18
19use crate::env::{Environment, SpaceInfo, SpaceType, StepInfo, StepResult};
20
21/// Snapshot of [`ContinuousLqr`]'s simulation state.
22///
23/// `ContinuousLqr` is fully deterministic (no internal RNG), so
24/// [`Environment::restore_state`] followed by [`Environment::step`] reproduces
25/// every subsequent [`StepResult`] bit-for-bit.
26#[derive(Debug, Clone)]
27pub struct ContinuousLqrState {
28    /// Position scalar.
29    pub position: f32,
30    /// Velocity scalar.
31    pub velocity: f32,
32    /// Step counter.
33    pub steps: usize,
34}
35
36/// Force clamp range applied to the 1D action input.
37const ACTION_CLAMP: f32 = 1.0;
38
39/// Per-step integration constant (treated as a unit timestep).
40const DT: f32 = 1.0;
41
42/// Mild linear damping on velocity each step.
43const VELOCITY_DAMPING: f32 = 0.1;
44
45/// Default episode length cap.
46const DEFAULT_MAX_STEPS: usize = 50;
47
48/// 1D linear quadratic regulator with a continuous-force action.
49#[derive(Debug, Clone)]
50pub struct ContinuousLqr {
51    /// Position scalar.
52    position: f32,
53
54    /// Velocity scalar.
55    velocity: f32,
56
57    /// Step counter for the current episode.
58    steps: usize,
59
60    /// Maximum number of steps before truncation.
61    max_steps: usize,
62}
63
64impl ContinuousLqr {
65    /// Create a new env with the default episode length.
66    pub fn new() -> Self {
67        Self::with_max_steps(DEFAULT_MAX_STEPS)
68    }
69
70    /// Create a new env with a custom maximum episode length.
71    pub fn with_max_steps(max_steps: usize) -> Self {
72        Self { position: 0.5, velocity: 0.0, steps: 0, max_steps }
73    }
74
75    /// Current position (for inspection / tests).
76    pub fn position(&self) -> f32 {
77        self.position
78    }
79
80    /// Current velocity (for inspection / tests).
81    pub fn velocity(&self) -> f32 {
82        self.velocity
83    }
84}
85
86impl Default for ContinuousLqr {
87    fn default() -> Self {
88        Self::new()
89    }
90}
91
92impl Environment for ContinuousLqr {
93    /// Continuous force action. Length-1 `Vec<f32>`; values outside
94    /// `[-ACTION_CLAMP, ACTION_CLAMP]` are clamped.
95    type Action = Vec<f32>;
96
97    /// Snapshot type. `ContinuousLqr` is deterministic (no RNG), so
98    /// restore + step reproduces subsequent results exactly.
99    type State = ContinuousLqrState;
100
101    fn reset(&mut self) {
102        self.position = 0.5;
103        self.velocity = 0.0;
104        self.steps = 0;
105    }
106
107    fn get_observation(&self) -> Vec<f32> {
108        vec![self.position, self.velocity]
109    }
110
111    fn step(&mut self, action: Vec<f32>) -> StepResult {
112        // Length-1 action; if a caller hands us an empty vec we
113        // treat it as zero force rather than panicking. Anything
114        // beyond index 0 is ignored.
115        let raw = action.first().copied().unwrap_or(0.0);
116        let force = raw.clamp(-ACTION_CLAMP, ACTION_CLAMP);
117
118        // Integrate: v_{t+1} = v_t + (force - damping*v_t)*dt
119        //            x_{t+1} = x_t + v_{t+1}*dt
120        self.velocity += (force - VELOCITY_DAMPING * self.velocity) * DT;
121        self.position += self.velocity * DT;
122
123        self.steps += 1;
124        let truncated = self.steps >= self.max_steps;
125
126        let reward = -(self.position * self.position) - 0.01 * (force * force);
127
128        StepResult {
129            observation: self.get_observation(),
130            reward,
131            terminated: false,
132            truncated,
133            info: StepInfo::default(),
134        }
135    }
136
137    fn observation_space(&self) -> SpaceInfo {
138        // [position, velocity], both continuous.
139        SpaceInfo { shape: vec![2], space_type: SpaceType::Box }
140    }
141
142    fn action_space(&self) -> SpaceInfo {
143        // 1D continuous force.
144        SpaceInfo { shape: vec![1], space_type: SpaceType::Box }
145    }
146
147    fn render(&self) -> Vec<u8> {
148        Vec::new()
149    }
150
151    fn close(&mut self) {}
152
153    fn clone_state(&self) -> ContinuousLqrState {
154        ContinuousLqrState { position: self.position, velocity: self.velocity, steps: self.steps }
155    }
156
157    fn restore_state(&mut self, state: &ContinuousLqrState) {
158        self.position = state.position;
159        self.velocity = state.velocity;
160        self.steps = state.steps;
161    }
162}
163
164#[cfg(test)]
165mod tests {
166    use super::*;
167
168    #[test]
169    fn continuous_lqr_accepts_vec_f32_action() {
170        let mut env = ContinuousLqr::new();
171        env.reset();
172
173        // Apply a positive force; velocity should become positive
174        // and position should move forward by the new velocity.
175        let result = env.step(vec![0.5]);
176        assert_eq!(result.observation.len(), 2);
177        assert!(env.velocity() > 0.0, "positive force should produce positive velocity");
178        assert!(env.position() > 0.5, "positive velocity should advance position");
179
180        // Reward should be negative (we are off the origin and used
181        // control effort).
182        assert!(result.reward < 0.0);
183
184        // Single step does not terminate or truncate the episode.
185        assert!(!result.terminated);
186        assert!(!result.truncated);
187    }
188
189    #[test]
190    fn continuous_lqr_terminates_after_max_steps() {
191        let mut env = ContinuousLqr::with_max_steps(3);
192        env.reset();
193
194        for _ in 0..2 {
195            let r = env.step(vec![0.0]);
196            assert!(!r.truncated);
197        }
198
199        let r = env.step(vec![0.0]);
200        assert!(r.truncated, "episode should truncate after max_steps");
201    }
202
203    #[test]
204    fn continuous_lqr_clamps_extreme_actions() {
205        let mut env = ContinuousLqr::new();
206        env.reset();
207
208        // Force way above clamp range: effective force is clamped
209        // to ACTION_CLAMP, so velocity change is bounded.
210        let _ = env.step(vec![1000.0]);
211        assert!(env.velocity() <= ACTION_CLAMP);
212    }
213
214    #[test]
215    fn continuous_lqr_action_space_is_box() {
216        let env = ContinuousLqr::new();
217        let space = env.action_space();
218        assert_eq!(space.shape, vec![1]);
219        assert!(matches!(space.space_type, SpaceType::Box));
220    }
221
222    #[test]
223    fn continuous_lqr_empty_action_treated_as_zero() {
224        let mut env = ContinuousLqr::new();
225        env.reset();
226        // No panic on empty action.
227        let _ = env.step(Vec::new());
228        assert_eq!(env.velocity(), -VELOCITY_DAMPING * 0.0);
229    }
230}
thrust_rl/env/games/continuous_lqr.rs

thrust_rl/env/games/
continuous_lqr.rs
