From af2ee7c7ce322c2f314e5a21be724a4c27863f67 Mon Sep 17 00:00:00 2001
From: Lars Bogner <upyxo@student.kit.edu>
Date: Thu, 30 Jul 2026 16:40:02 +0200
Subject: [PATCH] Add gumbel router configs sweeping learnable-knob
 combinations

Extends the 10-expert EnergyRouter + physical-conditioning benchmark
config with the new opt-in gumbel combine weights, isolating the
learn_centers/learn_temperature axis: none, centers only, and
centers+temperature.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
---
 ...er_energy_n10_gumbel_centers_physical.toml | 22 ++++++++++++++++++
 ...ergy_n10_gumbel_centers_temp_physical.toml | 23 +++++++++++++++++++
 .../router_energy_n10_gumbel_physical.toml    | 22 ++++++++++++++++++
 3 files changed, 67 insertions(+)
 create mode 100644 configs/router_energy_n10_gumbel_centers_physical.toml
 create mode 100644 configs/router_energy_n10_gumbel_centers_temp_physical.toml
 create mode 100644 configs/router_energy_n10_gumbel_physical.toml

diff --git a/configs/router_energy_n10_gumbel_centers_physical.toml b/configs/router_energy_n10_gumbel_centers_physical.toml
new file mode 100644
index 0000000..386e01e
--- /dev/null
+++ b/configs/router_energy_n10_gumbel_centers_physical.toml
@@ -0,0 +1,22 @@
+[train]
+mode          = "flow"
+epochs        = 30
+lr            = 3e-4
+warmup_epochs = 3
+val_fraction  = 0.1
+num_workers   = 4
+
+[model]
+conditioning = "physical"
+dropout      = 0.0
+
+[model.router]
+enabled           = true
+type              = "energy"
+n_experts         = 10
+expert_hidden_dim = 128
+expert_n_blocks   = 4
+temperature       = 0.05
+lambda_balance    = 0.035
+learn_centers     = true
+gumbel            = true
diff --git a/configs/router_energy_n10_gumbel_centers_temp_physical.toml b/configs/router_energy_n10_gumbel_centers_temp_physical.toml
new file mode 100644
index 0000000..933aa4b
--- /dev/null
+++ b/configs/router_energy_n10_gumbel_centers_temp_physical.toml
@@ -0,0 +1,23 @@
+[train]
+mode          = "flow"
+epochs        = 30
+lr            = 3e-4
+warmup_epochs = 3
+val_fraction  = 0.1
+num_workers   = 4
+
+[model]
+conditioning = "physical"
+dropout      = 0.0
+
+[model.router]
+enabled           = true
+type              = "energy"
+n_experts         = 10
+expert_hidden_dim = 128
+expert_n_blocks   = 4
+temperature       = 0.05
+lambda_balance    = 0.035
+learn_centers     = true
+learn_temperature = true
+gumbel            = true
diff --git a/configs/router_energy_n10_gumbel_physical.toml b/configs/router_energy_n10_gumbel_physical.toml
new file mode 100644
index 0000000..51c6a39
--- /dev/null
+++ b/configs/router_energy_n10_gumbel_physical.toml
@@ -0,0 +1,22 @@
+[train]
+mode          = "flow"
+epochs        = 30
+lr            = 3e-4
+warmup_epochs = 3
+val_fraction  = 0.1
+num_workers   = 4
+
+[model]
+conditioning = "physical"
+dropout      = 0.0
+
+[model.router]
+enabled           = true
+type              = "energy"
+n_experts         = 10
+expert_hidden_dim = 128
+expert_n_blocks   = 4
+temperature       = 0.05
+lambda_balance    = 0.035
+learn_centers     = false
+gumbel            = true