feat: upgrade notebook to Qwen3.5-4B with H100 hyperparams

- Model: unsloth/Qwen3.5-4B (from Qwen3-0.6B)
- num_generations: 8 (from 4) for smoother GRPO advantages
- learning_rate: 5e-5 (from 2e-4) for more stable training
- per_device_train_batch_size: 8 (from 4)

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

training/notebooks/fusion_design_lab_training.ipynb

@@ -5,14 +5,14 @@
    "id": "7fb27b941602401d91542211134fc71a",
    "metadata": {},
    "source": [
-    "# Fusion Design Lab \u2014 GRPO Training\n",
+    "# Fusion Design Lab — GRPO Training\n",
     "\n",
     "Train an LLM to optimize stellarator fusion reactor designs using **GRPO** (Group Relative Policy Optimization) with **Unsloth** and **TRL**.\n",
     "\n",
     "The agent interacts with a constrained optimization environment where it adjusts 4 geometric knobs of a stellarator boundary, aiming to **minimize max elongation** while satisfying 3 hard physics constraints:\n",
-    "- `aspect_ratio \u2264 4.0`\n",
-    "- `average_triangularity \u2264 -0.5`\n",
-    "- `edge_iota_over_nfp \u2265 0.3`\n",
+    "- `aspect_ratio ≤ 4.0`\n",
+    "- `average_triangularity ≤ -0.5`\n",
+    "- `edge_iota_over_nfp ≥ 0.3`\n",
     "\n",
     "Each episode has **6 evaluations** budgeted. The agent produces a plan of actions and the environment scores it via the `constellaration` physics verifier.\n",
     "\n",
@@ -59,29 +59,7 @@
    "id": "72eea5119410473aa328ad9291626812",
    "metadata": {},
    "outputs": [],
-   "source": [
-    "from unsloth import FastLanguageModel\n",
-    "\n",
-    "MODEL_NAME = \"unsloth/Qwen3-0.6B\"\n",
-    "MAX_SEQ_LENGTH = 2048\n",
-    "\n",
-    "model, tokenizer = FastLanguageModel.from_pretrained(\n",
-    "    model_name=MODEL_NAME,\n",
-    "    max_seq_length=MAX_SEQ_LENGTH,\n",
-    "    load_in_4bit=True,\n",
-    "    fast_inference=True,\n",
-    ")\n",
-    "\n",
-    "model = FastLanguageModel.get_peft_model(\n",
-    "    model,\n",
-    "    r=32,\n",
-    "    target_modules=[\"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\", \"gate_proj\", \"up_proj\", \"down_proj\"],\n",
-    "    lora_alpha=32,\n",
-    "    use_gradient_checkpointing=\"unsloth\",\n",
-    ")\n",
-    "\n",
-    "print(f\"Model loaded: {MODEL_NAME}\")"
-   ]
+   "source": "from unsloth import FastLanguageModel\n\nMODEL_NAME = \"unsloth/Qwen3.5-4B\"\nMAX_SEQ_LENGTH = 2048\n\nmodel, tokenizer = FastLanguageModel.from_pretrained(\n    model_name=MODEL_NAME,\n    max_seq_length=MAX_SEQ_LENGTH,\n    load_in_4bit=True,\n    fast_inference=True,\n)\n\nmodel = FastLanguageModel.get_peft_model(\n    model,\n    r=32,\n    target_modules=[\"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\", \"gate_proj\", \"up_proj\", \"down_proj\"],\n    lora_alpha=32,\n    use_gradient_checkpointing=\"unsloth\",\n)\n\nprint(f\"Model loaded: {MODEL_NAME}\")"
   },
   {
    "cell_type": "markdown",
@@ -220,7 +198,7 @@
    "source": [
     "## 6. Reward Function\n",
     "\n",
-    "The environment reward executes each generated action plan in the stellarator environment and returns the cumulative low-fidelity Reward V0 from the live environment. The environment's built-in reward decomposes feasibility (+3/-3 crossing bonuses, feasibility progress), objective (max elongation improvement), step costs, and failure penalties \u2014 see `server/environment.py:_compute_reward_breakdown(...)`.\n",
+    "The environment reward executes each generated action plan in the stellarator environment and returns the cumulative low-fidelity Reward V0 from the live environment. The environment's built-in reward decomposes feasibility (+3/-3 crossing bonuses, feasibility progress), objective (max elongation improvement), step costs, and failure penalties — see `server/environment.py:_compute_reward_breakdown(...)`.\n",
     "\n",
     "For the current training workflow, the notebook ignores `submit` and does not auto-submit. GRPO therefore optimizes the low-fidelity `run` path only. The live observation telemetry still exposes `reward_breakdown` and `action_monitor` for debugging reward behavior.\n"
    ]
@@ -313,41 +291,7 @@
    "id": "8a65eabff63a45729fe45fb5ade58bdc",
    "metadata": {},
    "outputs": [],
-   "source": [
-    "from trl import GRPOConfig, GRPOTrainer\n",
-    "\n",
-    "MAX_PROMPT_LENGTH = 768\n",
-    "MAX_COMPLETION_LENGTH = MAX_SEQ_LENGTH - MAX_PROMPT_LENGTH\n",
-    "\n",
-    "training_args = GRPOConfig(\n",
-    "    output_dir=\"./grpo_fusion_output\",\n",
-    "    learning_rate=2e-4,\n",
-    "    num_generations=4,\n",
-    "    max_completion_length=MAX_COMPLETION_LENGTH,\n",
-    "    max_prompt_length=MAX_PROMPT_LENGTH,\n",
-    "    per_device_train_batch_size=4,\n",
-    "    gradient_accumulation_steps=1,\n",
-    "    max_steps=60,\n",
-    "    temperature=1.0,\n",
-    "    logging_steps=1,\n",
-    "    save_steps=20,\n",
-    "    bf16=True,\n",
-    "    report_to=\"none\",\n",
-    "    seed=42,\n",
-    ")\n",
-    "\n",
-    "trainer = GRPOTrainer(\n",
-    "    model=model,\n",
-    "    processing_class=tokenizer,\n",
-    "    reward_funcs=[environment_reward_fn],\n",
-    "    args=training_args,\n",
-    "    train_dataset=dataset,\n",
-    ")\n",
-    "\n",
-    "print(\"Starting GRPO training...\")\n",
-    "train_result = trainer.train()\n",
-    "print(f\"Training complete. Total steps: {train_result.global_step}\")"
-   ]
+   "source": "from trl import GRPOConfig, GRPOTrainer\n\nMAX_PROMPT_LENGTH = 768\nMAX_COMPLETION_LENGTH = MAX_SEQ_LENGTH - MAX_PROMPT_LENGTH\n\ntraining_args = GRPOConfig(\n    output_dir=\"./grpo_fusion_output\",\n    learning_rate=5e-5,\n    num_generations=8,\n    max_completion_length=MAX_COMPLETION_LENGTH,\n    max_prompt_length=MAX_PROMPT_LENGTH,\n    per_device_train_batch_size=8,\n    gradient_accumulation_steps=1,\n    max_steps=60,\n    temperature=1.0,\n    logging_steps=1,\n    save_steps=20,\n    bf16=True,\n    report_to=\"none\",\n    seed=42,\n)\n\ntrainer = GRPOTrainer(\n    model=model,\n    processing_class=tokenizer,\n    reward_funcs=[environment_reward_fn],\n    args=training_args,\n    train_dataset=dataset,\n)\n\nprint(\"Starting GRPO training...\")\ntrain_result = trainer.train()\nprint(f\"Training complete. Total steps: {train_result.global_step}\")"
   },
   {
    "cell_type": "markdown",
@@ -401,7 +345,7 @@
     "else:\n",
     "    axes[1].text(0.5, 0.5, \"Reward metrics not logged\", ha=\"center\", va=\"center\")\n",
     "\n",
-    "plt.suptitle(\"Fusion Design Lab \u2014 GRPO Training Curves\", fontsize=14, fontweight=\"bold\")\n",
+    "plt.suptitle(\"Fusion Design Lab — GRPO Training Curves\", fontsize=14, fontweight=\"bold\")\n",
     "plt.tight_layout()\n",
     "plt.savefig(\"training_curves.png\", dpi=150, bbox_inches=\"tight\")\n",
     "plt.show()\n",
@@ -474,7 +418,7 @@
     "    episode = run_episode_with_actions(actions, seed_idx=seed_idx)\n",
     "    trace = [\n",
     "        (\n",
-    "            f\"{step.action_label} \u2192 reward={step.reward:.3f} \"\n",
+    "            f\"{step.action_label} → reward={step.reward:.3f} \"\n",
     "            f\"score={step.p1_score:.4f} feasible={step.constraints_satisfied} \"\n",
     "            f\"terms={reward_term_summary(step)}\"\n",
     "        )\n",
@@ -497,7 +441,7 @@
     "for seed in range(len(RESET_SEEDS)):\n",
     "    reward, trace = run_episode_with_model(seed)\n",
     "    trained_rewards.append(reward)\n",
-    "    print(f\"\\nSeed {seed} \u2014 Total reward: {reward:.3f}\")\n",
+    "    print(f\"\\nSeed {seed} — Total reward: {reward:.3f}\")\n",
     "    for line in trace:\n",
     "        print(f\"  {line}\")\n",
     "\n",
@@ -511,7 +455,7 @@
     "    seed_rewards = [run_random_episode(seed) for _ in range(10)]\n",
     "    random_rewards.extend(seed_rewards)\n",
     "    print(\n",
-    "        f\"Seed {seed} \u2014 Mean: {sum(seed_rewards) / len(seed_rewards):.3f}, Best: {max(seed_rewards):.3f}\"\n",
+    "        f\"Seed {seed} — Mean: {sum(seed_rewards) / len(seed_rewards):.3f}, Best: {max(seed_rewards):.3f}\"\n",
     "    )\n",
     "\n",
     "print(f\"\\nMean random reward: {sum(random_rewards) / len(random_rewards):.3f}\")\n",
@@ -550,7 +494,7 @@
     "with FusionLabClient(base_url=HF_SPACE_URL) as env:\n",
     "    reset_result = env.reset(seed=42)\n",
     "    remote_obs = reset_result.observation\n",
-    "    print(f\"\\nRemote reset \u2014 max_elongation: {remote_obs.max_elongation:.4f}\")\n",
+    "    print(f\"\\nRemote reset — max_elongation: {remote_obs.max_elongation:.4f}\")\n",
     "    print(f\"  aspect_ratio: {remote_obs.aspect_ratio:.4f}\")\n",
     "    print(f\"  constraints_satisfied: {remote_obs.constraints_satisfied}\")\n",
     "    print(f\"  budget_remaining: {remote_obs.budget_remaining}\")\n",
@@ -576,7 +520,7 @@
     "        print(\n",
     "            f\"  Step {i}: {action.intent} {action.parameter or ''} \"\n",
     "            f\"{action.direction or ''} {action.magnitude or ''} \"\n",
-    "            f\"\u2192 reward={reward:.3f}, score={step_obs.p1_score:.4f}, terms={reward_term_summary(step_obs)}\"\n",
+    "            f\"→ reward={reward:.3f}, score={step_obs.p1_score:.4f}, terms={reward_term_summary(step_obs)}\"\n",
     "        )\n",
     "        if result.done:\n",
     "            print(f\"  Episode done. Final score: {step_obs.p1_score:.4f}\")\n",