[Feat] Add decision dimension to builder and config

endure.toml

@@ -223,6 +223,7 @@ output_features = ["z0_cost", "z1_cost", "q_cost", "w_cost"]
 embedding_size = 8
 hidden_length = 3
 hidden_width = 32
+decision_dim = 64
 
 # Dropout percentage
 dropout = 0.0

endure/lcm/model/builder.py

@@ -46,6 +46,7 @@ def build_model(self, policy: Policy) -> torch.nn.Module:
             "hidden_length": model_params["hidden_length"],
             "hidden_width": model_params["hidden_width"],
             "dropout_percentage": model_params["dropout"],
+            "decision_dim": model_params["decision_dim"],
         }
 
         if model_params["norm_layer"] == "Batch":

endure/ltune/util/ltune_eval.py

@@ -6,7 +6,7 @@
 
 from endure.lcm.util import eval_lcm_impl
 from endure.lsm.cost import EndureCost
-from endure.lsm.types import LSMDesign, System, Policy
+from endure.lsm.types import LSMDesign, System, Policy, STR_POLICY_DICT
 from endure.ltune.data.generator import LTuneGenerator
 from endure.ltune.loss import LearnedCostModelLoss
 import endure.lsm.solver as Solver
@@ -19,6 +19,7 @@ def __init__(
         model: torch.nn.Module,
         design_type: str = "Level",
     ) -> None:
+        self.policy = STR_POLICY_DICT.get(design_type, Policy.KHybrid)
         self.gen = LTuneGenerator(config)
         self.loss = LearnedCostModelLoss(
             config,