thrust_rl/train/sac/
config.rs

1//! Soft Actor-Critic (SAC) configuration and hyperparameters.
2//!
3//! [`SacConfig`] mirrors the builder + `validate()` style of
4//! [`crate::train::dqn::DQNConfig`] and
5//! [`crate::train::ppo::PPOConfig`]. Defaults track Haarnoja et al. 2018
6//! v2 ("Soft Actor-Critic Algorithms and Applications", arXiv:1812.05905)
7//! on Pendulum-scale continuous-control tasks. This is the config half of
8//! PR E of the SAC decomposition (#136); the [`crate::train::sac::SacTrainer`]
9//! consumes it.
10
11use anyhow::{Result, anyhow};
12
13/// SAC configuration parameters.
14///
15/// Default values target classic Pendulum-scale continuous control:
16/// 1M-capacity replay, 256-sample batches, always-soft Polyak target
17/// updates (`tau = 0.005`), automatic entropy temperature tuning, and the
18/// `3e-4` Adam learning rate the v2 paper uses for all three optimizers.
19/// Smoke tests typically override `buffer_capacity` down to ~50k.
20#[derive(Debug, Clone)]
21pub struct SacConfig {
22    /// Adam learning rate for the stochastic actor.
23    pub actor_lr: f64,
24
25    /// Adam learning rate for both online critics (`q1`, `q2`).
26    pub critic_lr: f64,
27
28    /// Adam learning rate for the entropy temperature `log_alpha`. Only
29    /// used when [`Self::auto_alpha`] is `true`.
30    pub alpha_lr: f64,
31
32    /// Discount factor used in the critic TD target.
33    pub gamma: f64,
34
35    /// Polyak (soft) target update coefficient `tau`. SAC always performs
36    /// a soft update of both target critics every gradient step:
37    /// `theta_target <- tau * theta_online + (1 - tau) * theta_target`.
38    pub tau: f64,
39
40    /// Number of transitions sampled per gradient update.
41    pub batch_size: usize,
42
43    /// Maximum number of transitions stored in the replay buffer. Older
44    /// transitions are evicted FIFO once capacity is reached.
45    pub buffer_capacity: usize,
46
47    /// Minimum number of transitions required before the first gradient
48    /// update. Until the buffer holds this many transitions the trainer
49    /// only collects experience.
50    pub min_buffer_size: usize,
51
52    /// Number of environment steps for which actions are drawn uniformly
53    /// at random (pure exploration) before the actor starts choosing
54    /// actions.
55    pub learning_starts: usize,
56
57    /// Number of gradient updates performed per environment step (the
58    /// update-to-data ratio).
59    pub gradient_steps_per_env_step: usize,
60
61    /// Width of every hidden layer in the actor and critics.
62    pub hidden_dim: usize,
63
64    /// Number of hidden layers in the actor and critic trunks.
65    pub num_hidden_layers: usize,
66
67    /// Automatically tune the entropy temperature `alpha` (Haarnoja et al.
68    /// 2018 v2). When `true`, `log_alpha` is optimized to drive the
69    /// policy entropy toward [`Self::target_entropy`]. When `false`,
70    /// `alpha` is held fixed at [`Self::init_alpha`].
71    pub auto_alpha: bool,
72
73    /// Initial entropy temperature `alpha`. When [`Self::auto_alpha`] is
74    /// `false` this is the fixed value used throughout training; when
75    /// `true` it is the starting point (`log_alpha = ln(init_alpha)`).
76    pub init_alpha: f32,
77
78    /// Target policy entropy for automatic temperature tuning. `None`
79    /// resolves to the conventional heuristic `-action_dim` at trainer
80    /// construction time.
81    pub target_entropy: Option<f32>,
82
83    /// Optional global gradient-norm clip applied to every optimizer
84    /// step. `None` (the SAC default) leaves the updates unclipped.
85    pub max_grad_norm: Option<f64>,
86
87    /// Seed threaded through replay sampling, actor noise sampling, and
88    /// seeded network init for bit-exact reproducibility.
89    pub seed: u64,
90}
91
92impl Default for SacConfig {
93    fn default() -> Self {
94        Self {
95            actor_lr: 3e-4,
96            critic_lr: 3e-4,
97            alpha_lr: 3e-4,
98            gamma: 0.99,
99            tau: 0.005,
100            batch_size: 256,
101            buffer_capacity: 1_000_000,
102            min_buffer_size: 1_000,
103            learning_starts: 1_000,
104            gradient_steps_per_env_step: 1,
105            hidden_dim: 256,
106            num_hidden_layers: 2,
107            auto_alpha: true,
108            init_alpha: 0.2,
109            target_entropy: None,
110            max_grad_norm: None,
111            seed: 0,
112        }
113    }
114}
115
116impl SacConfig {
117    /// Create a new default configuration.
118    pub fn new() -> Self {
119        Self::default()
120    }
121
122    /// Validate configuration parameters.
123    ///
124    /// Returns an `Err` describing the first invalid field encountered.
125    pub fn validate(&self) -> Result<()> {
126        if self.actor_lr <= 0.0 {
127            return Err(anyhow!("actor_lr must be positive, got {}", self.actor_lr));
128        }
129        if self.critic_lr <= 0.0 {
130            return Err(anyhow!("critic_lr must be positive, got {}", self.critic_lr));
131        }
132        if self.alpha_lr <= 0.0 {
133            return Err(anyhow!("alpha_lr must be positive, got {}", self.alpha_lr));
134        }
135        if !(self.tau > 0.0 && self.tau <= 1.0) {
136            return Err(anyhow!("tau must be in (0, 1], got {}", self.tau));
137        }
138        if !(0.0..=1.0).contains(&self.gamma) {
139            return Err(anyhow!("gamma must be in [0, 1], got {}", self.gamma));
140        }
141        if self.batch_size == 0 {
142            return Err(anyhow!("batch_size must be positive"));
143        }
144        if self.buffer_capacity < self.batch_size {
145            return Err(anyhow!(
146                "buffer_capacity ({}) must be at least batch_size ({})",
147                self.buffer_capacity,
148                self.batch_size
149            ));
150        }
151        if self.min_buffer_size > self.buffer_capacity {
152            return Err(anyhow!(
153                "min_buffer_size ({}) must be <= buffer_capacity ({})",
154                self.min_buffer_size,
155                self.buffer_capacity
156            ));
157        }
158        if self.init_alpha <= 0.0 {
159            return Err(anyhow!("init_alpha must be positive, got {}", self.init_alpha));
160        }
161        if self.gradient_steps_per_env_step == 0 {
162            return Err(anyhow!("gradient_steps_per_env_step must be positive"));
163        }
164        if self.hidden_dim == 0 {
165            return Err(anyhow!("hidden_dim must be positive"));
166        }
167        if self.num_hidden_layers == 0 {
168            return Err(anyhow!("num_hidden_layers must be positive"));
169        }
170        if let Some(clip) = self.max_grad_norm
171            && clip <= 0.0
172        {
173            return Err(anyhow!("max_grad_norm must be positive when set, got {}", clip));
174        }
175        if let Some(te) = self.target_entropy
176            && !te.is_finite()
177        {
178            return Err(anyhow!("target_entropy must be finite when set, got {}", te));
179        }
180        Ok(())
181    }
182
183    /// Resolve the effective target entropy for automatic temperature
184    /// tuning, applying the `-action_dim` heuristic when
185    /// [`Self::target_entropy`] is `None`.
186    pub fn resolved_target_entropy(&self, action_dim: usize) -> f32 {
187        self.target_entropy.unwrap_or(-(action_dim as f32))
188    }
189
190    // ----- Builder-style setters (mirroring DQNConfig / PPOConfig) -----
191
192    /// Set the actor learning rate.
193    pub fn actor_lr(mut self, lr: f64) -> Self {
194        self.actor_lr = lr;
195        self
196    }
197
198    /// Set the critic learning rate (applied to both online critics).
199    pub fn critic_lr(mut self, lr: f64) -> Self {
200        self.critic_lr = lr;
201        self
202    }
203
204    /// Set the entropy-temperature learning rate.
205    pub fn alpha_lr(mut self, lr: f64) -> Self {
206        self.alpha_lr = lr;
207        self
208    }
209
210    /// Set the discount factor `gamma`.
211    pub fn gamma(mut self, gamma: f64) -> Self {
212        self.gamma = gamma;
213        self
214    }
215
216    /// Set the Polyak soft-update coefficient `tau`.
217    pub fn tau(mut self, tau: f64) -> Self {
218        self.tau = tau;
219        self
220    }
221
222    /// Set the minibatch size.
223    pub fn batch_size(mut self, size: usize) -> Self {
224        self.batch_size = size;
225        self
226    }
227
228    /// Set the replay buffer capacity.
229    pub fn buffer_capacity(mut self, capacity: usize) -> Self {
230        self.buffer_capacity = capacity;
231        self
232    }
233
234    /// Set the minimum buffer size before the first gradient update.
235    pub fn min_buffer_size(mut self, size: usize) -> Self {
236        self.min_buffer_size = size;
237        self
238    }
239
240    /// Set the number of random-action warmup steps.
241    pub fn learning_starts(mut self, steps: usize) -> Self {
242        self.learning_starts = steps;
243        self
244    }
245
246    /// Set the number of gradient updates per environment step.
247    pub fn gradient_steps_per_env_step(mut self, steps: usize) -> Self {
248        self.gradient_steps_per_env_step = steps;
249        self
250    }
251
252    /// Set the hidden-layer width for the actor and critics.
253    pub fn hidden_dim(mut self, dim: usize) -> Self {
254        self.hidden_dim = dim;
255        self
256    }
257
258    /// Set the number of hidden layers for the actor and critics.
259    pub fn num_hidden_layers(mut self, layers: usize) -> Self {
260        self.num_hidden_layers = layers;
261        self
262    }
263
264    /// Enable or disable automatic entropy-temperature tuning.
265    pub fn auto_alpha(mut self, enabled: bool) -> Self {
266        self.auto_alpha = enabled;
267        self
268    }
269
270    /// Set the initial (or fixed) entropy temperature `alpha`.
271    pub fn init_alpha(mut self, alpha: f32) -> Self {
272        self.init_alpha = alpha;
273        self
274    }
275
276    /// Set an explicit target entropy, overriding the `-action_dim`
277    /// heuristic.
278    pub fn target_entropy(mut self, entropy: f32) -> Self {
279        self.target_entropy = Some(entropy);
280        self
281    }
282
283    /// Enable global gradient-norm clipping with the given cap.
284    pub fn max_grad_norm(mut self, norm: f64) -> Self {
285        self.max_grad_norm = Some(norm);
286        self
287    }
288
289    /// Set the reproducibility seed.
290    pub fn seed(mut self, seed: u64) -> Self {
291        self.seed = seed;
292        self
293    }
294}
295
296#[cfg(test)]
297mod tests {
298    use super::*;
299
300    #[test]
301    fn test_default_config_validates() {
302        let cfg = SacConfig::default();
303        assert!(cfg.validate().is_ok());
304        assert_eq!(cfg.actor_lr, 3e-4);
305        assert_eq!(cfg.critic_lr, 3e-4);
306        assert_eq!(cfg.alpha_lr, 3e-4);
307        assert_eq!(cfg.gamma, 0.99);
308        assert_eq!(cfg.tau, 0.005);
309        assert_eq!(cfg.batch_size, 256);
310        assert_eq!(cfg.buffer_capacity, 1_000_000);
311        assert_eq!(cfg.min_buffer_size, 1_000);
312        assert_eq!(cfg.learning_starts, 1_000);
313        assert_eq!(cfg.gradient_steps_per_env_step, 1);
314        assert_eq!(cfg.hidden_dim, 256);
315        assert_eq!(cfg.num_hidden_layers, 2);
316        assert!(cfg.auto_alpha);
317        assert_eq!(cfg.init_alpha, 0.2);
318        assert!(cfg.target_entropy.is_none());
319        assert!(cfg.max_grad_norm.is_none());
320        assert_eq!(cfg.seed, 0);
321    }
322
323    #[test]
324    fn test_builder() {
325        let cfg = SacConfig::new()
326            .actor_lr(1e-3)
327            .critic_lr(2e-3)
328            .alpha_lr(5e-4)
329            .gamma(0.95)
330            .tau(0.01)
331            .batch_size(128)
332            .buffer_capacity(50_000)
333            .min_buffer_size(500)
334            .learning_starts(2_000)
335            .gradient_steps_per_env_step(2)
336            .hidden_dim(64)
337            .num_hidden_layers(3)
338            .auto_alpha(false)
339            .init_alpha(0.1)
340            .target_entropy(-2.0)
341            .max_grad_norm(5.0)
342            .seed(42);
343        assert!(cfg.validate().is_ok());
344        assert_eq!(cfg.actor_lr, 1e-3);
345        assert_eq!(cfg.critic_lr, 2e-3);
346        assert_eq!(cfg.alpha_lr, 5e-4);
347        assert_eq!(cfg.gamma, 0.95);
348        assert_eq!(cfg.tau, 0.01);
349        assert_eq!(cfg.batch_size, 128);
350        assert_eq!(cfg.buffer_capacity, 50_000);
351        assert_eq!(cfg.min_buffer_size, 500);
352        assert_eq!(cfg.learning_starts, 2_000);
353        assert_eq!(cfg.gradient_steps_per_env_step, 2);
354        assert_eq!(cfg.hidden_dim, 64);
355        assert_eq!(cfg.num_hidden_layers, 3);
356        assert!(!cfg.auto_alpha);
357        assert_eq!(cfg.init_alpha, 0.1);
358        assert_eq!(cfg.target_entropy, Some(-2.0));
359        assert_eq!(cfg.max_grad_norm, Some(5.0));
360        assert_eq!(cfg.seed, 42);
361    }
362
363    #[test]
364    fn test_resolved_target_entropy() {
365        let cfg = SacConfig::default();
366        assert_eq!(cfg.resolved_target_entropy(1), -1.0);
367        assert_eq!(cfg.resolved_target_entropy(3), -3.0);
368        let cfg = cfg.target_entropy(-7.5);
369        assert_eq!(cfg.resolved_target_entropy(3), -7.5);
370    }
371
372    #[test]
373    fn test_validate_rejects_non_positive_lrs() {
374        assert!(SacConfig::new().actor_lr(0.0).validate().is_err());
375        assert!(SacConfig::new().actor_lr(-1.0).validate().is_err());
376        assert!(SacConfig::new().critic_lr(0.0).validate().is_err());
377        assert!(SacConfig::new().alpha_lr(-1e-4).validate().is_err());
378    }
379
380    #[test]
381    fn test_validate_rejects_tau_out_of_range() {
382        assert!(SacConfig::new().tau(0.0).validate().is_err());
383        assert!(SacConfig::new().tau(-0.1).validate().is_err());
384        assert!(SacConfig::new().tau(1.5).validate().is_err());
385        assert!(SacConfig::new().tau(1.0).validate().is_ok());
386        assert!(SacConfig::new().tau(0.005).validate().is_ok());
387    }
388
389    #[test]
390    fn test_validate_rejects_gamma_out_of_range() {
391        assert!(SacConfig::new().gamma(-0.1).validate().is_err());
392        assert!(SacConfig::new().gamma(1.5).validate().is_err());
393        assert!(SacConfig::new().gamma(0.0).validate().is_ok());
394        assert!(SacConfig::new().gamma(1.0).validate().is_ok());
395    }
396
397    #[test]
398    fn test_validate_rejects_zero_batch() {
399        assert!(SacConfig::new().batch_size(0).validate().is_err());
400    }
401
402    #[test]
403    fn test_validate_rejects_capacity_below_batch() {
404        let cfg = SacConfig::new().buffer_capacity(64).batch_size(256);
405        assert!(cfg.validate().is_err());
406    }
407
408    #[test]
409    fn test_validate_rejects_min_buffer_above_capacity() {
410        let cfg = SacConfig::new().buffer_capacity(1_000).min_buffer_size(5_000);
411        assert!(cfg.validate().is_err());
412    }
413
414    #[test]
415    fn test_validate_rejects_non_positive_init_alpha() {
416        assert!(SacConfig::new().init_alpha(0.0).validate().is_err());
417        assert!(SacConfig::new().init_alpha(-0.2).validate().is_err());
418        assert!(SacConfig::new().init_alpha(0.2).validate().is_ok());
419    }
420
421    #[test]
422    fn test_validate_rejects_zero_gradient_steps() {
423        assert!(SacConfig::new().gradient_steps_per_env_step(0).validate().is_err());
424    }
425
426    #[test]
427    fn test_validate_rejects_zero_hidden_and_layers() {
428        assert!(SacConfig::new().hidden_dim(0).validate().is_err());
429        assert!(SacConfig::new().num_hidden_layers(0).validate().is_err());
430    }
431
432    #[test]
433    fn test_validate_rejects_non_positive_grad_clip() {
434        assert!(SacConfig::new().max_grad_norm(0.0).validate().is_err());
435        assert!(SacConfig::new().max_grad_norm(-1.0).validate().is_err());
436        assert!(SacConfig::new().max_grad_norm(10.0).validate().is_ok());
437    }
438}
thrust_rl/train/sac/config.rs

thrust_rl/train/sac/
config.rs
