Fix convergence bugs in REINFORCE and Q-Learning

REINFORCE:
- Replace broken global-baseline approach with per-episode return
  normalization: G̃_t = (G_t − μ) / (σ + ε). The old baseline used
  the running average of G_0 for ALL timesteps t, causing massive
  negative advantages at late steps (G_{T-1}≈1, baseline≈300 → −299)
  and a zero gradient at t=0 on the first episode.
- Lower default lr 0.005 → 0.002 (research: >0.003 often diverges)
- Lower default gamma 0.97 → 0.99 (0.97 too myopic for 500-step horizon)

Q-Learning:
- Add alpha (learning rate) decay per episode alongside epsilon decay.
  Constant alpha prevents Q-values from ever stabilizing — they keep
  jittering regardless of how many episodes run.
- Raise initial alpha default 0.2 → 0.5 (decay brings it to 0.1 min)
- Raise epsilonMin default 0.01 → 0.05 (0.01 stops exploration too early)
- Tighten epsilonDecay default 0.998 → 0.995 (faster initial exploration)

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: wadekarg

src/algorithms/classicCartpole/discretizedQLearning.ts

@@ -8,33 +8,39 @@ import { argmax, randInt } from '../../utils/math'
  * Q-Learning over discretized classic CartPole state space.
  * Bins the 4 continuous state variables (x, xDot, theta, thetaDot)
  * into discrete buckets. Default: 6×6×12×12 = 5,184 states.
+ *
+ * Both α and ε decay per episode — constant α prevents Q-values from stabilizing.
  */
 export class DiscretizedQLearningAgent implements Agent<ClassicCartPoleState, ClassicCartPoleAction> {
   private qTable: Map<string, number[]> = new Map()
-  private alpha: number
+  private alpha: number          // initial learning rate
   private gamma: number
-  private epsilon: number
+  private epsilon: number        // initial exploration rate
   private epsilonDecay: number
   private epsilonMin: number
   private currentEpsilon: number
+  private currentAlpha: number   // decayed learning rate
+  private readonly alphaDecay = 0.999
+  private readonly alphaMin = 0.1
   private episodeCount = 0
   private discretizationConfig: ClassicDiscretizationConfig
   private numActions = 2
 
   constructor(
-    alpha = 0.1,
+    alpha = 0.5,
     gamma = 0.99,
-    epsilon = 0.1,
+    epsilon = 1.0,
     discretizationConfig: ClassicDiscretizationConfig = DEFAULT_CLASSIC_DISCRETIZATION,
     epsilonDecay = 0.995,
-    epsilonMin = 0.01,
+    epsilonMin = 0.05,
   ) {
     this.alpha = alpha
     this.gamma = gamma
     this.epsilon = epsilon
     this.epsilonDecay = epsilonDecay
     this.epsilonMin = epsilonMin
     this.currentEpsilon = epsilon
+    this.currentAlpha = alpha
     this.discretizationConfig = discretizationConfig
   }
 
@@ -60,14 +66,20 @@ export class DiscretizedQLearningAgent implements Agent<ClassicCartPoleState, Cl
     const nextQ = this.getQ(nextKey)
 
     const target = done ? reward : reward + this.gamma * Math.max(...nextQ)
-    q[action] += this.alpha * (target - q[action])
+    q[action] += this.currentAlpha * (target - q[action])
 
     if (done) {
       this.episodeCount++
+      // Decay ε per episode
       this.currentEpsilon = Math.max(
         this.epsilonMin,
         this.epsilon * Math.pow(this.epsilonDecay, this.episodeCount),
       )
+      // Decay α per episode — constant α prevents convergence
+      this.currentAlpha = Math.max(
+        this.alphaMin,
+        this.alpha * Math.pow(this.alphaDecay, this.episodeCount),
+      )
     }
   }
 
@@ -87,6 +99,7 @@ export class DiscretizedQLearningAgent implements Agent<ClassicCartPoleState, Cl
     this.qTable.clear()
     this.episodeCount = 0
     this.currentEpsilon = this.epsilon
+    this.currentAlpha = this.alpha
   }
 
   setParams(alpha: number, gamma: number, epsilon: number): void {

src/algorithms/classicCartpole/reinforce.ts

@@ -7,6 +7,12 @@ import type { ClassicCartPoleState, ClassicCartPoleAction } from '../../environm
  * Policy: π(a|s) = softmax(W · φ(s))
  * Features φ(s): [1, x/2.4, v/3.0, θ/0.21, ω/3.5, (θ/0.21)², (ω/3.5)²]  (7 features)
  * Weights: W is 2×7 matrix (2 actions × 7 features)
+ *
+ * Return normalization (per-episode):
+ *   G̃_t = (G_t − mean(G)) / (std(G) + ε)
+ *   This is critical: using the same global baseline for all timesteps creates high variance
+ *   because G_t naturally decreases across the episode (G_T-1 ≈ 1, G_0 ≈ T).
+ *   Per-episode normalization removes this structural bias.
  */
 
 interface Transition {
@@ -46,10 +52,12 @@ export class ReinforceAgent implements Agent<ClassicCartPoleState, ClassicCartPo
   private lr: number
   private gamma: number
   private trajectory: Transition[] = []
-  private baselineReturn = 0
+
+  // Display-only baseline: running average of episode returns (not used for learning)
+  private displayBaseline = 0
   private episodeCount = 0
 
-  constructor(lr = 0.01, gamma = 0.99) {
+  constructor(lr = 0.002, gamma = 0.99) {
     this.lr = lr
     this.gamma = gamma
     this.weights = Array.from({ length: NUM_ACTIONS }, () =>
@@ -76,29 +84,37 @@ export class ReinforceAgent implements Agent<ClassicCartPoleState, ClassicCartPo
     const T = this.trajectory.length
     const returns: number[] = new Array(T)
 
+    // Reward-to-go: G_t = r_t + γ·r_{t+1} + γ²·r_{t+2} + …
     let G = 0
     for (let t = T - 1; t >= 0; t--) {
       G = this.trajectory[t].reward + this.gamma * G
       returns[t] = G
     }
 
+    // Per-episode return normalization: G̃_t = (G_t − μ) / (σ + ε)
+    // Fixes the high-variance problem caused by G_t naturally shrinking across the episode.
+    const meanG = returns.reduce((a, b) => a + b, 0) / T
+    const varG = returns.reduce((a, v) => a + (v - meanG) ** 2, 0) / T
+    const stdG = Math.sqrt(varG + 1e-8)
+
+    // Update display baseline (running average of total episode return — for UI only)
     this.episodeCount++
-    const episodeReturn = returns[0]
-    this.baselineReturn += (episodeReturn - this.baselineReturn) / this.episodeCount
+    this.displayBaseline += (returns[0] - this.displayBaseline) / this.episodeCount
 
     for (let t = 0; t < T; t++) {
       const { state: s, action: a } = this.trajectory[t]
-      const advantage = returns[t] - this.baselineReturn
+      const normalizedReturn = (returns[t] - meanG) / stdG
 
       const phi = features(s)
       const logits = this.weights.map((w) =>
         w.reduce((sum, wi, i) => sum + wi * phi[i], 0),
       )
       const probs = softmax(logits)
 
+      // Policy gradient: ∇W = lr · G̃_t · (I(j==a) − π(j|s)) · φ(s)
       for (let j = 0; j < NUM_ACTIONS; j++) {
         const indicator = j === a ? 1 : 0
-        const gradScale = (indicator - probs[j]) * advantage
+        const gradScale = (indicator - probs[j]) * normalizedReturn
         for (let f = 0; f < NUM_FEATURES; f++) {
           this.weights[j][f] += this.lr * gradScale * phi[f]
         }
@@ -110,15 +126,15 @@ export class ReinforceAgent implements Agent<ClassicCartPoleState, ClassicCartPo
 
   getValues(): Record<string, number[]> {
     const flat = this.weights.flat()
-    return { weights: flat, baseline: [this.baselineReturn] }
+    return { weights: flat, baseline: [this.displayBaseline] }
   }
 
   reset(): void {
     this.weights = Array.from({ length: NUM_ACTIONS }, () =>
       new Array(NUM_FEATURES).fill(0),
     )
     this.trajectory = []
-    this.baselineReturn = 0
+    this.displayBaseline = 0
     this.episodeCount = 0
   }
 

src/components/classicCartpole/ClassicCartPolePage.tsx

@@ -24,14 +24,14 @@ const ALGO_CATEGORY: Record<AlgorithmType, string> = {
 
 export function ClassicCartPolePage() {
   const [algorithmType, setAlgorithmType] = useState<AlgorithmType>('discretized-q')
-  const [alpha, setAlpha] = useState(0.2)
+  const [alpha, setAlpha] = useState(0.5)
   const [gamma, setGamma] = useState(0.99)
-  const [gammaRF, setGammaRF] = useState(0.97)
+  const [gammaRF, setGammaRF] = useState(0.99)
   const [epsilon, setEpsilon] = useState(1.0)
-  const [lr, setLr] = useState(0.005)
+  const [lr, setLr] = useState(0.002)
   const [bins, setBins] = useState(6)
-  const [epsilonDecay, setEpsilonDecay] = useState(0.998)
-  const [epsilonMin, setEpsilonMin] = useState(0.01)
+  const [epsilonDecay, setEpsilonDecay] = useState(0.995)
+  const [epsilonMin, setEpsilonMin] = useState(0.05)
   const [showIntro, setShowIntro] = useState(true)
   const [maxSteps, setMaxSteps] = useState(100000)
   const [envSeed, setEnvSeed] = useState(0)