From 332b3eb1208cd892169cade6ce9ab7af98eed89e Mon Sep 17 00:00:00 2001
From: jgrusewski <jgrusewski@localhost>
Date: Sun, 29 Mar 2026 18:16:47 +0200
Subject: [PATCH] feat(training): cosine LR decay in TOML profile,
 min_epochs_before_stopping=80
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- Added lr_decay_type (0=constant, 1=linear, 2=cosine) and lr_min to
  the TOML training profile system. Wired through apply_to() with
  total_steps derived from epochs.

- dqn-localdev.toml: lr_decay_type=2 (cosine 1e-4→1e-5 over 200 epochs),
  min_epochs_before_stopping=80 (was 50, gives model more room to recover
  after C51 transition).

- Smoketest config unchanged (lr_decay_type not set → default Constant).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 config/training/dqn-localdev.toml |  4 +++-
 crates/ml/src/training_profile.rs | 19 +++++++++++++++++++
 2 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/config/training/dqn-localdev.toml b/config/training/dqn-localdev.toml
index 3e2833f5c..401cab730 100644
--- a/config/training/dqn-localdev.toml
+++ b/config/training/dqn-localdev.toml
@@ -15,6 +15,8 @@ hidden_dim_base = 64
 max_steps_per_epoch = 200
 reward_scale = 1.0
 huber_delta = 1.0
+lr_decay_type = 2
+lr_min = 0.00001
 
 [distributional]
 num_atoms = 51
@@ -35,7 +37,7 @@ min_replay_size = 500
 [early_stopping]
 enabled = true
 patience = 20
-min_epochs_before_stopping = 50
+min_epochs_before_stopping = 80
 
 [experience]
 gpu_n_episodes = 64
diff --git a/crates/ml/src/training_profile.rs b/crates/ml/src/training_profile.rs
index 945d94f10..5a0bafd5f 100644
--- a/crates/ml/src/training_profile.rs
+++ b/crates/ml/src/training_profile.rs
@@ -71,6 +71,10 @@ pub struct TrainingSection {
     pub huber_delta: Option<f64>,
     /// Adam epsilon — BF16 requires ≥1e-4 (1e-8 rounds to 0 → div-by-zero).
     pub adam_epsilon: Option<f64>,
+    /// LR decay type: 0=constant, 1=linear, 2=cosine (default: 0).
+    pub lr_decay_type: Option<u32>,
+    /// Minimum LR for cosine/linear decay (default: 1e-6).
+    pub lr_min: Option<f64>,
 }
 
 /// Epsilon-greedy exploration parameters (DQN-specific).
@@ -653,6 +657,21 @@ impl DqnTrainingProfile {
             if let Some(v) = t.adam_epsilon {
                 hp.adam_epsilon = v;
             }
+            if let Some(decay) = t.lr_decay_type {
+                let total_steps = t.epochs.unwrap_or(hp.epochs);
+                let min_lr = t.lr_min.unwrap_or(1e-6);
+                hp.lr_decay_type = match decay {
+                    1 => crate::trainers::dqn::lr_scheduler::LRDecayType::Linear {
+                        end_lr: min_lr,
+                        total_steps,
+                    },
+                    2 => crate::trainers::dqn::lr_scheduler::LRDecayType::Cosine {
+                        min_lr,
+                        total_steps,
+                    },
+                    _ => crate::trainers::dqn::lr_scheduler::LRDecayType::Constant,
+                };
+            }
         }
 
         // [exploration]