fix: correct GRPO training hyperparameters to prevent KL explosion

beta=0.01 weakened the KL penalty (opposite of intended), causing policy
to diverge at step 18 and collapse to learn×28. Fix: beta=0.1 (stronger
constraint), lr=5e-5 (more conservative), max_grad_norm=0.5 (clipping).

Also fix train.py: max_completion_length was 368 (prompt-pad remainder)
instead of 16, which would allow verbose drift in standalone runs.

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

training/RhythmEnv_GRPO_Training.ipynb

@@ -224,7 +224,51 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": "from trl import GRPOConfig, GRPOTrainer\n\nMAX_STEPS = 500  # Increase to 1000 if time allows\nNUM_GENERATIONS = 4\nLEARNING_RATE = 2e-4\n\nmax_prompt_length = 400\nmax_completion_length = 16  # Action names are 3-15 chars — no need for more\n\ntraining_args = GRPOConfig(\n    temperature=1.0,\n    learning_rate=LEARNING_RATE,\n    beta=0.01,                  # KL penalty coefficient (called beta in TRL, default 0.04 causes KL explosion)\n    weight_decay=0.001,\n    warmup_ratio=0.1,\n    lr_scheduler_type=\"linear\",\n    optim=\"adamw_8bit\",\n    logging_steps=1,\n    per_device_train_batch_size=1,\n    gradient_accumulation_steps=4,\n    num_generations=NUM_GENERATIONS,\n    max_prompt_length=max_prompt_length,\n    max_completion_length=max_completion_length,\n    max_steps=MAX_STEPS,\n    save_steps=100,\n    report_to=REPORT_TO,\n    output_dir=\"outputs/rhythmenv_trained\",\n)\n\ntrainer = GRPOTrainer(\n    model=model,\n    processing_class=tokenizer,\n    reward_funcs=reward_funcs,\n    args=training_args,\n    train_dataset=dataset,\n)\n\nprint(f\"Training config: {MAX_STEPS} steps, {NUM_GENERATIONS} generations, lr={LEARNING_RATE}\")\nprint(f\"  beta=0.01 (KL penalty — reduced from default to prevent policy drift)\")\nprint(f\"  max_completion_length=16 (action names only, no verbose outputs)\")\nprint(\"Starting training...\")"
+   "source": [
+    "from trl import GRPOConfig, GRPOTrainer\n",
+    "\n",
+    "MAX_STEPS = 500  # Increase to 1000 if time allows\n",
+    "NUM_GENERATIONS = 4\n",
+    "LEARNING_RATE = 5e-5  # Reduced from default — lower lr prevents destabilizing early gradient steps\n",
+    "\n",
+    "max_prompt_length = 400\n",
+    "max_completion_length = 16  # Action names are 3-15 chars\n",
+    "\n",
+    "training_args = GRPOConfig(\n",
+    "    temperature=1.0,\n",
+    "    learning_rate=LEARNING_RATE,\n",
+    "    beta=0.1,                   # KL penalty — higher = more conservative, prevents policy drift\n",
+    "    max_grad_norm=0.5,          # Gradient clipping prevents large destabilizing updates\n",
+    "    weight_decay=0.001,\n",
+    "    warmup_ratio=0.1,\n",
+    "    lr_scheduler_type=\"linear\",\n",
+    "    optim=\"adamw_8bit\",\n",
+    "    logging_steps=1,\n",
+    "    per_device_train_batch_size=1,\n",
+    "    gradient_accumulation_steps=4,\n",
+    "    num_generations=NUM_GENERATIONS,\n",
+    "    max_prompt_length=max_prompt_length,\n",
+    "    max_completion_length=max_completion_length,\n",
+    "    max_steps=MAX_STEPS,\n",
+    "    save_steps=100,\n",
+    "    report_to=REPORT_TO,\n",
+    "    output_dir=\"outputs/rhythmenv_trained\",\n",
+    ")\n",
+    "\n",
+    "trainer = GRPOTrainer(\n",
+    "    model=model,\n",
+    "    processing_class=tokenizer,\n",
+    "    reward_funcs=reward_funcs,\n",
+    "    args=training_args,\n",
+    "    train_dataset=dataset,\n",
+    ")\n",
+    "\n",
+    "print(f\"Training config: {MAX_STEPS} steps, {NUM_GENERATIONS} generations, lr={LEARNING_RATE}\")\n",
+    "print(f\"  beta=0.1 (higher KL penalty = more conservative = less policy drift)\")\n",
+    "print(f\"  max_grad_norm=0.5 (gradient clipping for stability)\")\n",
+    "print(f\"  max_completion_length=16 (action names only, no verbose outputs)\")\n",
+    "print(\"Starting training...\")"
+   ]
   },
   {
    "cell_type": "code",

training/train.py

@@ -32,7 +32,7 @@ def main():
                         help="Maximum training samples")
     parser.add_argument("--num_generations", type=int, default=4,
                         help="Number of completions per prompt for GRPO")
-    parser.add_argument("--learning_rate", type=float, default=2e-4,
+    parser.add_argument("--learning_rate", type=float, default=5e-5,
                         help="Learning rate")
     parser.add_argument("--output_dir", type=str, default="outputs/rhythmenv_trained",
                         help="Output directory for model and logs")
@@ -128,11 +128,13 @@ def main():
     from trl import GRPOConfig, GRPOTrainer
 
     max_prompt_length = 400
-    max_completion_length = max_seq_length - max_prompt_length
+    max_completion_length = 16  # Action names are 3-15 chars; cap prevents verbose drift
 
     training_args = GRPOConfig(
         temperature=1.0,
         learning_rate=args.learning_rate,
+        beta=0.1,               # KL penalty — higher = more conservative, prevents policy drift
+        max_grad_norm=0.5,      # Gradient clipping prevents large destabilizing updates
         weight_decay=0.001,
         warmup_ratio=0.1,
         lr_scheduler_type="linear",