thrust_rl/train/ppo/
config.rs

1//! PPO configuration and hyperparameters
2//!
3//! This module defines the configuration parameters for PPO training
4//! and provides validation and builder pattern methods.
5
6use anyhow::{Result, anyhow};
7
8/// PPO configuration parameters
9///
10/// These hyperparameters control the PPO training process.
11/// Default values are based on common settings that work well
12/// for simple environments like CartPole.
13#[derive(Debug, Clone)]
14pub struct PPOConfig {
15    /// Learning rate for policy and value function
16    pub learning_rate: f64,
17
18    /// Number of training epochs per rollout
19    pub n_epochs: usize,
20
21    /// Minibatch size for training
22    pub batch_size: usize,
23
24    /// Discount factor (gamma)
25    pub gamma: f64,
26
27    /// GAE lambda parameter
28    pub gae_lambda: f64,
29
30    /// PPO clipping parameter (epsilon)
31    pub clip_range: f64,
32
33    /// Value function clipping parameter
34    pub clip_range_vf: f64,
35
36    /// Value function loss coefficient
37    pub vf_coef: f64,
38
39    /// Entropy bonus coefficient
40    pub ent_coef: f64,
41
42    /// Maximum gradient norm for clipping.
43    ///
44    /// Applied by both
45    /// [`PPOTrainerBurn`](crate::train::ppo::trainer::PPOTrainerBurn) and
46    /// [`RecurrentPPOTrainer`](crate::train::ppo::recurrent_trainer::RecurrentPPOTrainer)
47    /// as a **global** L2-norm clip (see
48    /// [`clip_grads_by_global_norm`](crate::train::grad_clip::clip_grads_by_global_norm))
49    /// over the concatenation of every parameter's gradient, immediately
50    /// before each optimizer step (issue #299; prior to that fix the field
51    /// was silently ignored by both trainers). Must be positive —
52    /// [`validate`](PPOConfig::validate) rejects non-positive values — so to
53    /// effectively disable clipping set a very large cap (e.g. `1e9`).
54    pub max_grad_norm: f64,
55
56    /// Target KL divergence for early stopping
57    pub target_kl: f64,
58
59    /// Seed for the inner-loop RNG used to shuffle minibatch indices.
60    ///
61    /// Plumbed through the trainer so that two trainers built with the
62    /// same `seed` produce bit-identical updates on the same rollout
63    /// (issue #109). The PPO trainer constructs an internal
64    /// `StdRng::seed_from_u64(seed)` at `new()` time and uses it for
65    /// every minibatch shuffle.
66    pub seed: u64,
67}
68
69impl Default for PPOConfig {
70    fn default() -> Self {
71        Self {
72            learning_rate: 3e-4,
73            n_epochs: 10,
74            batch_size: 64,
75            gamma: 0.99,
76            gae_lambda: 0.95,
77            clip_range: 0.2,
78            clip_range_vf: 0.2,
79            vf_coef: 0.5,
80            ent_coef: 0.01,
81            max_grad_norm: 0.5,
82            target_kl: 0.01,
83            seed: 0,
84        }
85    }
86}
87
88impl PPOConfig {
89    /// Create a new default configuration
90    pub fn new() -> Self {
91        Self::default()
92    }
93
94    /// Validate configuration parameters
95    pub fn validate(&self) -> Result<()> {
96        if self.learning_rate <= 0.0 {
97            return Err(anyhow!("learning_rate must be positive"));
98        }
99        if self.n_epochs == 0 {
100            return Err(anyhow!("n_epochs must be positive"));
101        }
102        if self.batch_size == 0 {
103            return Err(anyhow!("batch_size must be positive"));
104        }
105        if !(0.0..=1.0).contains(&self.gamma) {
106            return Err(anyhow!("gamma must be in [0, 1]"));
107        }
108        if !(0.0..=1.0).contains(&self.gae_lambda) {
109            return Err(anyhow!("gae_lambda must be in [0, 1]"));
110        }
111        if self.clip_range <= 0.0 {
112            return Err(anyhow!("clip_range must be positive"));
113        }
114        if self.clip_range_vf <= 0.0 {
115            return Err(anyhow!("clip_range_vf must be positive"));
116        }
117        if self.vf_coef < 0.0 {
118            return Err(anyhow!("vf_coef must be non-negative"));
119        }
120        if self.ent_coef < 0.0 {
121            return Err(anyhow!("ent_coef must be non-negative"));
122        }
123        if self.max_grad_norm <= 0.0 {
124            return Err(anyhow!("max_grad_norm must be positive"));
125        }
126        if self.target_kl < 0.0 {
127            return Err(anyhow!("target_kl must be non-negative"));
128        }
129        Ok(())
130    }
131
132    /// Set learning rate
133    pub fn learning_rate(mut self, lr: f64) -> Self {
134        self.learning_rate = lr;
135        self
136    }
137
138    /// Set number of training epochs
139    pub fn n_epochs(mut self, epochs: usize) -> Self {
140        self.n_epochs = epochs;
141        self
142    }
143
144    /// Set minibatch size
145    pub fn batch_size(mut self, size: usize) -> Self {
146        self.batch_size = size;
147        self
148    }
149
150    /// Set discount factor
151    pub fn gamma(mut self, gamma: f64) -> Self {
152        self.gamma = gamma;
153        self
154    }
155
156    /// Set GAE lambda
157    pub fn gae_lambda(mut self, lambda: f64) -> Self {
158        self.gae_lambda = lambda;
159        self
160    }
161
162    /// Set PPO clipping parameter
163    pub fn clip_range(mut self, clip: f64) -> Self {
164        self.clip_range = clip;
165        self
166    }
167
168    /// Set value function clipping parameter
169    pub fn clip_range_vf(mut self, clip: f64) -> Self {
170        self.clip_range_vf = clip;
171        self
172    }
173
174    /// Set value function loss coefficient
175    pub fn vf_coef(mut self, coef: f64) -> Self {
176        self.vf_coef = coef;
177        self
178    }
179
180    /// Set entropy bonus coefficient
181    pub fn ent_coef(mut self, coef: f64) -> Self {
182        self.ent_coef = coef;
183        self
184    }
185
186    /// Set maximum gradient norm
187    pub fn max_grad_norm(mut self, norm: f64) -> Self {
188        self.max_grad_norm = norm;
189        self
190    }
191
192    /// Set target KL divergence
193    pub fn target_kl(mut self, kl: f64) -> Self {
194        self.target_kl = kl;
195        self
196    }
197
198    /// Set the seed for the inner-loop RNG used for minibatch shuffling.
199    ///
200    /// See [`PPOConfig::seed`].
201    pub fn seed(mut self, seed: u64) -> Self {
202        self.seed = seed;
203        self
204    }
205}
206
207#[cfg(test)]
208mod tests {
209    use super::*;
210
211    #[test]
212    fn test_default_config() {
213        let config = PPOConfig::default();
214        assert!(config.validate().is_ok());
215        assert_eq!(config.learning_rate, 3e-4);
216        assert_eq!(config.n_epochs, 10);
217        assert_eq!(config.batch_size, 64);
218    }
219
220    #[test]
221    fn test_config_validation() {
222        // Valid config should pass
223        let config = PPOConfig::new();
224        assert!(config.validate().is_ok());
225
226        // Invalid learning rate
227        let config = PPOConfig::new().learning_rate(-1.0);
228        assert!(config.validate().is_err());
229
230        // Invalid gamma
231        let config = PPOConfig::new().gamma(1.5);
232        assert!(config.validate().is_err());
233
234        // Invalid n_epochs
235        let config = PPOConfig::new().n_epochs(0);
236        assert!(config.validate().is_err());
237
238        // Invalid batch_size
239        let config = PPOConfig::new().batch_size(0);
240        assert!(config.validate().is_err());
241
242        // Invalid clip_range
243        let config = PPOConfig::new().clip_range(-0.1);
244        assert!(config.validate().is_err());
245
246        // Invalid vf_coef (should allow 0.0)
247        let config = PPOConfig::new().vf_coef(-0.1);
248        assert!(config.validate().is_err());
249
250        // Valid vf_coef = 0.0
251        let config = PPOConfig::new().vf_coef(0.0);
252        assert!(config.validate().is_ok());
253    }
254
255    #[test]
256    fn test_config_builder() {
257        let config = PPOConfig::new()
258            .learning_rate(1e-3)
259            .n_epochs(5)
260            .batch_size(128)
261            .gamma(0.95)
262            .clip_range(0.1);
263
264        assert_eq!(config.learning_rate, 1e-3);
265        assert_eq!(config.n_epochs, 5);
266        assert_eq!(config.batch_size, 128);
267        assert_eq!(config.gamma, 0.95);
268        assert_eq!(config.clip_range, 0.1);
269
270        // Other values should remain default
271        assert_eq!(config.gae_lambda, 0.95);
272        assert_eq!(config.vf_coef, 0.5);
273    }
274}
thrust_rl/train/ppo/config.rs

thrust_rl/train/ppo/
config.rs
