thrust_rl/train/a2c/
config.rs

1//! A2C configuration and hyperparameters.
2//!
3//! [`A2cConfig`] mirrors the builder + `validate()` style of
4//! [`crate::train::ppo::PPOConfig`] and [`crate::train::sac::SacConfig`].
5//! Defaults track classic synchronous Advantage Actor-Critic (Mnih et al.
6//! 2016, "Asynchronous Methods for Deep Reinforcement Learning",
7//! arXiv:1602.01783) as popularized by the synchronous A2C variant: a
8//! `7e-4` learning rate, 5-step rollouts across 16 parallel actors, and
9//! full n-step returns (`gae_lambda = 1.0`). This is the config half of
10//! PR A of the A2C decomposition (#150); the A2C trainer (#152) consumes
11//! it.
12
13use anyhow::{Result, anyhow};
14
15/// A2C configuration parameters.
16///
17/// These hyperparameters control the synchronous Advantage Actor-Critic
18/// training process. Default values target classic discrete-control tasks
19/// like CartPole: a single gradient update per `n_steps`-long rollout
20/// collected across `num_envs` synchronous actors, full n-step returns
21/// (`gae_lambda = 1.0`), and global gradient-norm clipping.
22#[derive(Debug, Clone)]
23pub struct A2cConfig {
24    /// Learning rate for the shared actor-critic optimizer.
25    pub learning_rate: f64,
26
27    /// Discount factor (`gamma`).
28    pub gamma: f64,
29
30    /// GAE lambda parameter. The classic A2C default of `1.0` recovers
31    /// full n-step returns (no GAE smoothing).
32    pub gae_lambda: f64,
33
34    /// Value function loss coefficient.
35    pub value_coef: f64,
36
37    /// Entropy bonus coefficient.
38    pub entropy_coef: f64,
39
40    /// Rollout length: number of environment steps collected per update
41    /// (Mnih et al. 2016 use 5-step rollouts).
42    pub n_steps: usize,
43
44    /// Number of parallel synchronous actors.
45    pub num_envs: usize,
46
47    /// Maximum global gradient norm for clipping.
48    pub max_grad_norm: f64,
49
50    /// Whether to normalize advantages to zero-mean/unit-variance per
51    /// rollout before computing the policy loss.
52    pub normalize_advantages: bool,
53
54    /// Seed threaded into seeded network init (e.g.
55    /// [`MlpBurnConfig::with_seed`](crate::policy::mlp::MlpBurnConfig::with_seed))
56    /// for reproducibility. Two trainers built with the same `seed`
57    /// produce bit-identical initialization.
58    pub seed: u64,
59}
60
61impl Default for A2cConfig {
62    fn default() -> Self {
63        Self {
64            learning_rate: 7e-4,
65            gamma: 0.99,
66            gae_lambda: 1.0,
67            value_coef: 0.5,
68            entropy_coef: 0.01,
69            n_steps: 5,
70            num_envs: 16,
71            max_grad_norm: 0.5,
72            normalize_advantages: true,
73            seed: 0,
74        }
75    }
76}
77
78impl A2cConfig {
79    /// Create a new default configuration.
80    pub fn new() -> Self {
81        Self::default()
82    }
83
84    /// Validate configuration parameters.
85    ///
86    /// Returns an `Err` describing the first invalid field encountered.
87    pub fn validate(&self) -> Result<()> {
88        if self.learning_rate <= 0.0 {
89            return Err(anyhow!("learning_rate must be positive, got {}", self.learning_rate));
90        }
91        if !(0.0..=1.0).contains(&self.gamma) {
92            return Err(anyhow!("gamma must be in [0, 1], got {}", self.gamma));
93        }
94        if !(0.0..=1.0).contains(&self.gae_lambda) {
95            return Err(anyhow!("gae_lambda must be in [0, 1], got {}", self.gae_lambda));
96        }
97        if self.value_coef < 0.0 {
98            return Err(anyhow!("value_coef must be non-negative, got {}", self.value_coef));
99        }
100        if self.entropy_coef < 0.0 {
101            return Err(anyhow!("entropy_coef must be non-negative, got {}", self.entropy_coef));
102        }
103        if self.n_steps == 0 {
104            return Err(anyhow!("n_steps must be positive"));
105        }
106        if self.num_envs == 0 {
107            return Err(anyhow!("num_envs must be positive"));
108        }
109        if self.max_grad_norm <= 0.0 {
110            return Err(anyhow!("max_grad_norm must be positive, got {}", self.max_grad_norm));
111        }
112        Ok(())
113    }
114
115    // ----- Builder-style setters (mirroring PPOConfig / SacConfig) -----
116
117    /// Set the learning rate.
118    pub fn learning_rate(mut self, lr: f64) -> Self {
119        self.learning_rate = lr;
120        self
121    }
122
123    /// Set the discount factor `gamma`.
124    pub fn gamma(mut self, gamma: f64) -> Self {
125        self.gamma = gamma;
126        self
127    }
128
129    /// Set the GAE lambda parameter.
130    pub fn gae_lambda(mut self, lambda: f64) -> Self {
131        self.gae_lambda = lambda;
132        self
133    }
134
135    /// Set the value function loss coefficient.
136    pub fn value_coef(mut self, coef: f64) -> Self {
137        self.value_coef = coef;
138        self
139    }
140
141    /// Set the entropy bonus coefficient.
142    pub fn entropy_coef(mut self, coef: f64) -> Self {
143        self.entropy_coef = coef;
144        self
145    }
146
147    /// Set the rollout length (steps collected per update).
148    pub fn n_steps(mut self, steps: usize) -> Self {
149        self.n_steps = steps;
150        self
151    }
152
153    /// Set the number of parallel synchronous actors.
154    pub fn num_envs(mut self, envs: usize) -> Self {
155        self.num_envs = envs;
156        self
157    }
158
159    /// Set the maximum global gradient norm for clipping.
160    pub fn max_grad_norm(mut self, norm: f64) -> Self {
161        self.max_grad_norm = norm;
162        self
163    }
164
165    /// Enable or disable per-rollout advantage normalization.
166    pub fn normalize_advantages(mut self, enabled: bool) -> Self {
167        self.normalize_advantages = enabled;
168        self
169    }
170
171    /// Set the reproducibility seed.
172    pub fn seed(mut self, seed: u64) -> Self {
173        self.seed = seed;
174        self
175    }
176}
177
178#[cfg(test)]
179mod tests {
180    use super::*;
181
182    #[test]
183    fn test_default_config() {
184        let config = A2cConfig::default();
185        assert!(config.validate().is_ok());
186        assert_eq!(config.learning_rate, 7e-4);
187        assert_eq!(config.gamma, 0.99);
188        assert_eq!(config.gae_lambda, 1.0);
189        assert_eq!(config.value_coef, 0.5);
190        assert_eq!(config.entropy_coef, 0.01);
191        assert_eq!(config.n_steps, 5);
192        assert_eq!(config.num_envs, 16);
193        assert_eq!(config.max_grad_norm, 0.5);
194        assert!(config.normalize_advantages);
195        assert_eq!(config.seed, 0);
196    }
197
198    #[test]
199    fn test_config_validation() {
200        // Valid config should pass
201        let config = A2cConfig::new();
202        assert!(config.validate().is_ok());
203
204        // Invalid learning rate
205        assert!(A2cConfig::new().learning_rate(0.0).validate().is_err());
206        assert!(A2cConfig::new().learning_rate(-1.0).validate().is_err());
207
208        // Invalid gamma
209        assert!(A2cConfig::new().gamma(-0.1).validate().is_err());
210        assert!(A2cConfig::new().gamma(1.5).validate().is_err());
211        assert!(A2cConfig::new().gamma(0.0).validate().is_ok());
212        assert!(A2cConfig::new().gamma(1.0).validate().is_ok());
213
214        // Invalid gae_lambda
215        assert!(A2cConfig::new().gae_lambda(-0.1).validate().is_err());
216        assert!(A2cConfig::new().gae_lambda(1.5).validate().is_err());
217        assert!(A2cConfig::new().gae_lambda(0.0).validate().is_ok());
218        assert!(A2cConfig::new().gae_lambda(1.0).validate().is_ok());
219
220        // Invalid value_coef (negative rejected, 0.0 accepted edge)
221        assert!(A2cConfig::new().value_coef(-0.1).validate().is_err());
222        assert!(A2cConfig::new().value_coef(0.0).validate().is_ok());
223
224        // Invalid entropy_coef (negative rejected, 0.0 accepted)
225        assert!(A2cConfig::new().entropy_coef(-0.1).validate().is_err());
226        assert!(A2cConfig::new().entropy_coef(0.0).validate().is_ok());
227
228        // Invalid n_steps
229        assert!(A2cConfig::new().n_steps(0).validate().is_err());
230
231        // Invalid num_envs
232        assert!(A2cConfig::new().num_envs(0).validate().is_err());
233
234        // Invalid max_grad_norm
235        assert!(A2cConfig::new().max_grad_norm(0.0).validate().is_err());
236        assert!(A2cConfig::new().max_grad_norm(-1.0).validate().is_err());
237    }
238
239    #[test]
240    fn test_config_builder() {
241        let config = A2cConfig::new()
242            .learning_rate(1e-3)
243            .gamma(0.95)
244            .gae_lambda(0.9)
245            .value_coef(0.25)
246            .entropy_coef(0.05)
247            .n_steps(20)
248            .num_envs(8)
249            .max_grad_norm(1.0)
250            .normalize_advantages(false)
251            .seed(42);
252
253        assert!(config.validate().is_ok());
254        assert_eq!(config.learning_rate, 1e-3);
255        assert_eq!(config.gamma, 0.95);
256        assert_eq!(config.gae_lambda, 0.9);
257        assert_eq!(config.value_coef, 0.25);
258        assert_eq!(config.entropy_coef, 0.05);
259        assert_eq!(config.n_steps, 20);
260        assert_eq!(config.num_envs, 8);
261        assert_eq!(config.max_grad_norm, 1.0);
262        assert!(!config.normalize_advantages);
263        assert_eq!(config.seed, 42);
264    }
265}
thrust_rl/train/a2c/config.rs

thrust_rl/train/a2c/
config.rs
