tune: GRPO hyperparameter fixes from ML reviewer

Root cause of likely under-training in v1: 4-layer reward stack with
format/legal saturating at +1 each was drowning the env_reward and
belief_accuracy meta-RL signals. Plus several TRL/DeepSeek best-practice
deviations.

Changes:
- beta: 0.1 -> 0.04 (TRL/DeepSeek default; was anchoring policy too hard)
- num_generations: 4 -> 8 (halves advantage-estimate variance in
continuous-profile setting; critical for GRPO signal)
- max_completion_length: 20 -> 32 (was silently truncating belief digits
for actions like FAMILY_TIME and BINGE_WATCH after BPE tokenization)
- reward_weights=[0.3, 0.3, 1.0, 1.0] in GRPOConfig: scales format_valid +
action_legal down so env_reward and belief_accuracy dominate the gradient.
Wrapped in try/except for TRL versions that don't support reward_weights.

FAST_MODE preset also bumped NUM_GENERATIONS from 2 -> 4. Two completions
per group is too low to estimate group advantage at all.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

scripts/train_on_hf.py

@@ -57,12 +57,13 @@ PLOTS_DIR = "/tmp/rhythm_env/plots"
 FAST_MODE = os.environ.get("FAST_MODE", "0") == "1"
 
 if FAST_MODE:
+    # Smoke-train preset: enough signal in 200 steps to decide go/no-go
     DEFAULTS = dict(MAX_STEPS=200, NUM_EPISODES=80, MAX_SAMPLES=800,
-                    NUM_GENERATIONS=2, LORA_RANK=8, BETA=0.1,
+                    NUM_GENERATIONS=4, LORA_RANK=8, BETA=0.04,
                     LEARNING_RATE=5e-5, EVAL_EPISODES=2)
 else:
     DEFAULTS = dict(MAX_STEPS=1500, NUM_EPISODES=300, MAX_SAMPLES=3000,
-                    NUM_GENERATIONS=4, LORA_RANK=8, BETA=0.1,
+                    NUM_GENERATIONS=8, LORA_RANK=8, BETA=0.04,
                     LEARNING_RATE=5e-5, EVAL_EPISODES=5)
 
 MAX_STEPS = int(os.environ.get("MAX_STEPS", str(DEFAULTS["MAX_STEPS"])))

training/train.py

@@ -34,11 +34,11 @@ def main():
                         help="Number of episodes for dataset generation (more diversity = better meta-RL)")
     parser.add_argument("--max_samples", type=int, default=3000,
                         help="Maximum training samples")
-    parser.add_argument("--num_generations", type=int, default=4,
-                        help="Completions per prompt for GRPO (higher = lower variance, more compute)")
+    parser.add_argument("--num_generations", type=int, default=8,
+                        help="Completions per prompt for GRPO (8 default, lower variance for continuous-profile meta-RL)")
     parser.add_argument("--learning_rate", type=float, default=5e-5)
-    parser.add_argument("--beta", type=float, default=0.1,
-                        help="KL penalty (raise to 0.2 if training is unstable)")
+    parser.add_argument("--beta", type=float, default=0.04,
+                        help="KL penalty (TRL/DeepSeek default; raise to 0.1+ if KL diverges)")
     parser.add_argument("--lora_rank", type=int, default=8,
                         help="LoRA rank (8 = more capacity than original 4 for meta-RL)")
     parser.add_argument("--hint_fraction", type=float, default=0.15,
@@ -141,10 +141,15 @@ def main():
 
     from trl import GRPOConfig, GRPOTrainer
 
-    max_prompt_length = 600          # bumped from 400 for longer prompts (history + hint)
-    max_completion_length = 20       # bumped from 16 for "ACTION_NAME D D D" format
+    max_prompt_length = 600          # history + hint room
+    max_completion_length = 32       # bumped from 20 to prevent silent truncation of belief digits
 
-    training_args = GRPOConfig(
+    # reward_weights: scale per-layer to prevent format/legal (saturated near +1) from
+    # drowning out env_reward and belief_accuracy (the actual learning signals).
+    # Order MUST match reward_funcs in main(): format_valid, action_legal, env_reward, belief_accuracy
+    reward_weights = [0.3, 0.3, 1.0, 1.0]
+
+    training_args_kwargs = dict(
         temperature=1.0,
         learning_rate=args.learning_rate,
         beta=args.beta,
@@ -160,10 +165,17 @@ def main():
         max_prompt_length=max_prompt_length,
         max_completion_length=max_completion_length,
         max_steps=args.max_steps,
-        save_steps=250,                 # checkpoint every 250 (was 100)
+        save_steps=250,
         report_to=args.report_to,
         output_dir=args.output_dir,
     )
+    # reward_weights was added in TRL 0.13+; pass only if supported
+    try:
+        training_args = GRPOConfig(**training_args_kwargs, reward_weights=reward_weights)
+        print(f"Using GRPOConfig with reward_weights={reward_weights}")
+    except TypeError:
+        training_args = GRPOConfig(**training_args_kwargs)
+        print("WARN: TRL version does not support reward_weights; using uniform weighting")
 
     print(f"max_steps={args.max_steps}, num_generations={args.num_generations}, "
           f"lr={args.learning_rate}, beta={args.beta}")