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    /// If true, replace GAE with V-trace (Espeholt et al. 2018) for
55    /// off-policy correction. Requires that target-policy log-probs be
56    /// supplied to the rollout preprocessing step. Defaults to `false`,
57    /// leaving existing A2C behavior unchanged.
58    pub use_vtrace: bool,
59
60    /// V-trace rho clipping threshold (Espeholt 2018, eq. 1). Ignored when
61    /// `use_vtrace = false`. Typical value: `1.0` (the IMPALA paper
62    /// default).
63    pub vtrace_rho_bar: f32,
64
65    /// V-trace c clipping threshold (Espeholt 2018, eq. 2). Ignored when
66    /// `use_vtrace = false`. Typical value: `1.0`.
67    pub vtrace_c_bar: f32,
68
69    /// Seed threaded into seeded network init (e.g.
70    /// [`MlpBurnConfig::with_seed`](crate::policy::mlp::MlpBurnConfig::with_seed))
71    /// for reproducibility. Two trainers built with the same `seed`
72    /// produce bit-identical initialization.
73    pub seed: u64,
74}
75
76impl Default for A2cConfig {
77    fn default() -> Self {
78        Self {
79            learning_rate: 7e-4,
80            gamma: 0.99,
81            gae_lambda: 1.0,
82            value_coef: 0.5,
83            entropy_coef: 0.01,
84            n_steps: 5,
85            num_envs: 16,
86            max_grad_norm: 0.5,
87            normalize_advantages: true,
88            use_vtrace: false,
89            vtrace_rho_bar: 1.0,
90            vtrace_c_bar: 1.0,
91            seed: 0,
92        }
93    }
94}
95
96impl A2cConfig {
97    /// Create a new default configuration.
98    pub fn new() -> Self {
99        Self::default()
100    }
101
102    /// Validate configuration parameters.
103    ///
104    /// Returns an `Err` describing the first invalid field encountered.
105    pub fn validate(&self) -> Result<()> {
106        if self.learning_rate <= 0.0 {
107            return Err(anyhow!("learning_rate must be positive, got {}", self.learning_rate));
108        }
109        if !(0.0..=1.0).contains(&self.gamma) {
110            return Err(anyhow!("gamma must be in [0, 1], got {}", self.gamma));
111        }
112        if !(0.0..=1.0).contains(&self.gae_lambda) {
113            return Err(anyhow!("gae_lambda must be in [0, 1], got {}", self.gae_lambda));
114        }
115        if self.value_coef < 0.0 {
116            return Err(anyhow!("value_coef must be non-negative, got {}", self.value_coef));
117        }
118        if self.entropy_coef < 0.0 {
119            return Err(anyhow!("entropy_coef must be non-negative, got {}", self.entropy_coef));
120        }
121        if self.n_steps == 0 {
122            return Err(anyhow!("n_steps must be positive"));
123        }
124        if self.num_envs == 0 {
125            return Err(anyhow!("num_envs must be positive"));
126        }
127        if self.max_grad_norm <= 0.0 {
128            return Err(anyhow!("max_grad_norm must be positive, got {}", self.max_grad_norm));
129        }
130        if self.vtrace_rho_bar <= 0.0 {
131            return Err(anyhow!("vtrace_rho_bar must be positive, got {}", self.vtrace_rho_bar));
132        }
133        if self.vtrace_c_bar <= 0.0 {
134            return Err(anyhow!("vtrace_c_bar must be positive, got {}", self.vtrace_c_bar));
135        }
136        Ok(())
137    }
138
139    // ----- Builder-style setters (mirroring PPOConfig / SacConfig) -----
140
141    /// Set the learning rate.
142    pub fn learning_rate(mut self, lr: f64) -> Self {
143        self.learning_rate = lr;
144        self
145    }
146
147    /// Set the discount factor `gamma`.
148    pub fn gamma(mut self, gamma: f64) -> Self {
149        self.gamma = gamma;
150        self
151    }
152
153    /// Set the GAE lambda parameter.
154    pub fn gae_lambda(mut self, lambda: f64) -> Self {
155        self.gae_lambda = lambda;
156        self
157    }
158
159    /// Set the value function loss coefficient.
160    pub fn value_coef(mut self, coef: f64) -> Self {
161        self.value_coef = coef;
162        self
163    }
164
165    /// Set the entropy bonus coefficient.
166    pub fn entropy_coef(mut self, coef: f64) -> Self {
167        self.entropy_coef = coef;
168        self
169    }
170
171    /// Set the rollout length (steps collected per update).
172    pub fn n_steps(mut self, steps: usize) -> Self {
173        self.n_steps = steps;
174        self
175    }
176
177    /// Set the number of parallel synchronous actors.
178    pub fn num_envs(mut self, envs: usize) -> Self {
179        self.num_envs = envs;
180        self
181    }
182
183    /// Set the maximum global gradient norm for clipping.
184    pub fn max_grad_norm(mut self, norm: f64) -> Self {
185        self.max_grad_norm = norm;
186        self
187    }
188
189    /// Enable or disable per-rollout advantage normalization.
190    pub fn normalize_advantages(mut self, enabled: bool) -> Self {
191        self.normalize_advantages = enabled;
192        self
193    }
194
195    /// Enable or disable V-trace off-policy correction (replaces GAE).
196    pub fn use_vtrace(mut self, enabled: bool) -> Self {
197        self.use_vtrace = enabled;
198        self
199    }
200
201    /// Set the V-trace rho clipping threshold (Espeholt 2018, eq. 1).
202    pub fn vtrace_rho_bar(mut self, rho_bar: f32) -> Self {
203        self.vtrace_rho_bar = rho_bar;
204        self
205    }
206
207    /// Set the V-trace c clipping threshold (Espeholt 2018, eq. 2).
208    pub fn vtrace_c_bar(mut self, c_bar: f32) -> Self {
209        self.vtrace_c_bar = c_bar;
210        self
211    }
212
213    /// Set the reproducibility seed.
214    pub fn seed(mut self, seed: u64) -> Self {
215        self.seed = seed;
216        self
217    }
218}
219
220#[cfg(test)]
221mod tests {
222    use super::*;
223
224    #[test]
225    fn test_default_config() {
226        let config = A2cConfig::default();
227        assert!(config.validate().is_ok());
228        assert_eq!(config.learning_rate, 7e-4);
229        assert_eq!(config.gamma, 0.99);
230        assert_eq!(config.gae_lambda, 1.0);
231        assert_eq!(config.value_coef, 0.5);
232        assert_eq!(config.entropy_coef, 0.01);
233        assert_eq!(config.n_steps, 5);
234        assert_eq!(config.num_envs, 16);
235        assert_eq!(config.max_grad_norm, 0.5);
236        assert!(config.normalize_advantages);
237        // V-trace defaults leave existing behavior unchanged.
238        assert!(!config.use_vtrace);
239        assert_eq!(config.vtrace_rho_bar, 1.0);
240        assert_eq!(config.vtrace_c_bar, 1.0);
241        assert_eq!(config.seed, 0);
242    }
243
244    #[test]
245    fn test_config_validation() {
246        // Valid config should pass
247        let config = A2cConfig::new();
248        assert!(config.validate().is_ok());
249
250        // Invalid learning rate
251        assert!(A2cConfig::new().learning_rate(0.0).validate().is_err());
252        assert!(A2cConfig::new().learning_rate(-1.0).validate().is_err());
253
254        // Invalid gamma
255        assert!(A2cConfig::new().gamma(-0.1).validate().is_err());
256        assert!(A2cConfig::new().gamma(1.5).validate().is_err());
257        assert!(A2cConfig::new().gamma(0.0).validate().is_ok());
258        assert!(A2cConfig::new().gamma(1.0).validate().is_ok());
259
260        // Invalid gae_lambda
261        assert!(A2cConfig::new().gae_lambda(-0.1).validate().is_err());
262        assert!(A2cConfig::new().gae_lambda(1.5).validate().is_err());
263        assert!(A2cConfig::new().gae_lambda(0.0).validate().is_ok());
264        assert!(A2cConfig::new().gae_lambda(1.0).validate().is_ok());
265
266        // Invalid value_coef (negative rejected, 0.0 accepted edge)
267        assert!(A2cConfig::new().value_coef(-0.1).validate().is_err());
268        assert!(A2cConfig::new().value_coef(0.0).validate().is_ok());
269
270        // Invalid entropy_coef (negative rejected, 0.0 accepted)
271        assert!(A2cConfig::new().entropy_coef(-0.1).validate().is_err());
272        assert!(A2cConfig::new().entropy_coef(0.0).validate().is_ok());
273
274        // Invalid n_steps
275        assert!(A2cConfig::new().n_steps(0).validate().is_err());
276
277        // Invalid num_envs
278        assert!(A2cConfig::new().num_envs(0).validate().is_err());
279
280        // Invalid max_grad_norm
281        assert!(A2cConfig::new().max_grad_norm(0.0).validate().is_err());
282        assert!(A2cConfig::new().max_grad_norm(-1.0).validate().is_err());
283
284        // V-trace clip thresholds must be positive.
285        assert!(A2cConfig::new().vtrace_rho_bar(0.0).validate().is_err());
286        assert!(A2cConfig::new().vtrace_rho_bar(-0.5).validate().is_err());
287        assert!(A2cConfig::new().vtrace_rho_bar(1.0).validate().is_ok());
288        assert!(A2cConfig::new().vtrace_c_bar(0.0).validate().is_err());
289        assert!(A2cConfig::new().vtrace_c_bar(-0.5).validate().is_err());
290        assert!(A2cConfig::new().vtrace_c_bar(1.0).validate().is_ok());
291        // Enabling V-trace with valid thresholds is accepted.
292        assert!(
293            A2cConfig::new()
294                .use_vtrace(true)
295                .vtrace_rho_bar(0.8)
296                .vtrace_c_bar(1.2)
297                .validate()
298                .is_ok()
299        );
300    }
301
302    #[test]
303    fn test_config_builder() {
304        let config = A2cConfig::new()
305            .learning_rate(1e-3)
306            .gamma(0.95)
307            .gae_lambda(0.9)
308            .value_coef(0.25)
309            .entropy_coef(0.05)
310            .n_steps(20)
311            .num_envs(8)
312            .max_grad_norm(1.0)
313            .normalize_advantages(false)
314            .use_vtrace(true)
315            .vtrace_rho_bar(0.9)
316            .vtrace_c_bar(1.1)
317            .seed(42);
318
319        assert!(config.validate().is_ok());
320        assert_eq!(config.learning_rate, 1e-3);
321        assert_eq!(config.gamma, 0.95);
322        assert_eq!(config.gae_lambda, 0.9);
323        assert_eq!(config.value_coef, 0.25);
324        assert_eq!(config.entropy_coef, 0.05);
325        assert_eq!(config.n_steps, 20);
326        assert_eq!(config.num_envs, 8);
327        assert_eq!(config.max_grad_norm, 1.0);
328        assert!(!config.normalize_advantages);
329        assert!(config.use_vtrace);
330        assert_eq!(config.vtrace_rho_bar, 0.9);
331        assert_eq!(config.vtrace_c_bar, 1.1);
332        assert_eq!(config.seed, 42);
333    }
334}
thrust_rl/train/a2c/config.rs

thrust_rl/train/a2c/
config.rs
