fix(reforce): optuna search space definition

Signed-off-by: Jérôme Benoit <jerome.benoit@piment-noir.org>

diff --git a/ReforceXY/user_data/freqaimodels/ReforceXY.py b/ReforceXY/user_data/freqaimodels/ReforceXY.py
index 47b2cfa13714e7b49212d3baa18c657428793662..d24f31b953b515f098146dc768c1b528d3687510 100644 (file)
--- a/ReforceXY/user_data/freqaimodels/ReforceXY.py
+++ b/ReforceXY/user_data/freqaimodels/ReforceXY.py
@@ -157,6 +157,7 @@ class ReforceXY(BaseReinforcementLearningModel):
         "QRDQN",
     )
     _SCHEDULE_TYPES: Final[tuple[ScheduleType, ...]] = ("linear", "constant", "unknown")
+    _SCHEDULE_TYPES_KNOWN: Final[tuple[ScheduleTypeKnown, ...]] = ("linear", "constant")
     _EXIT_POTENTIAL_MODES: Final[tuple[ExitPotentialMode, ...]] = (
         "canonical",
         "non_canonical",
@@ -186,6 +187,7 @@ class ReforceXY(BaseReinforcementLearningModel):
         "leaky_relu",
     )
     _OPTIMIZER_CLASSES: Final[tuple[OptimizerClass, ...]] = ("adam", "adamw", "rmsprop")
+    _OPTIMIZER_CLASSES_OPTUNA: Final[tuple[OptimizerClass, ...]] = ("adamw", "rmsprop")
     _NET_ARCH_SIZES: Final[tuple[NetArchSize, ...]] = (
         "small",
         "medium",
@@ -3997,10 +3999,10 @@ def get_common_ppo_optuna_params(trial: Trial) -> Dict[str, Any]:
         "max_grad_norm": trial.suggest_float("max_grad_norm", 0.3, 1.0, step=0.05),
         "vf_coef": trial.suggest_float("vf_coef", 0.0, 1.0, step=0.05),
         "lr_schedule": trial.suggest_categorical(
-            "lr_schedule", list(ReforceXY._SCHEDULE_TYPES)
+            "lr_schedule", list(ReforceXY._SCHEDULE_TYPES_KNOWN)
         ),
         "cr_schedule": trial.suggest_categorical(
-            "cr_schedule", list(ReforceXY._SCHEDULE_TYPES)
+            "cr_schedule", list(ReforceXY._SCHEDULE_TYPES_KNOWN)
         ),
         "target_kl": trial.suggest_categorical(
             "target_kl", [None, 0.01, 0.015, 0.02, 0.03, 0.04]
@@ -4013,8 +4015,7 @@ def get_common_ppo_optuna_params(trial: Trial) -> Dict[str, Any]:
             "activation_fn", list(ReforceXY._ACTIVATION_FUNCTIONS)
         ),
         "optimizer_class": trial.suggest_categorical(
-            "optimizer_class",
-            [ReforceXY._OPTIMIZER_CLASSES[1], ReforceXY._OPTIMIZER_CLASSES[2]],
+            "optimizer_class", list(ReforceXY._OPTIMIZER_CLASSES_OPTUNA)
         ),
     }
 
@@ -4070,8 +4071,8 @@ def get_common_dqn_optuna_params(trial: Trial) -> Dict[str, Any]:
         ),
         "learning_rate": trial.suggest_float("learning_rate", 1e-5, 3e-3, log=True),
         "lr_schedule": trial.suggest_categorical(
-            "lr_schedule", list(ReforceXY._SCHEDULE_TYPES[:2])
-        ),  # ["linear", "constant"]
+            "lr_schedule", list(ReforceXY._SCHEDULE_TYPES_KNOWN)
+        ),
         "buffer_size": trial.suggest_categorical(
             "buffer_size", [int(1e4), int(5e4), int(1e5), int(2e5)]
         ),
@@ -4093,8 +4094,7 @@ def get_common_dqn_optuna_params(trial: Trial) -> Dict[str, Any]:
             "activation_fn", list(ReforceXY._ACTIVATION_FUNCTIONS)
         ),
         "optimizer_class": trial.suggest_categorical(
-            "optimizer_class",
-            [ReforceXY._OPTIMIZER_CLASSES[1], ReforceXY._OPTIMIZER_CLASSES[2]],
+            "optimizer_class", list(ReforceXY._OPTIMIZER_CLASSES_OPTUNA)
         ),
     }
 

diff --git a/quickadapter/user_data/strategies/Utils.py b/quickadapter/user_data/strategies/Utils.py
index 7230382143511fde5a6c36bb8adaafb9e314a45a..f444d7ae41e0f45ec7d80724df98192e075c83fd 100644 (file)
--- a/quickadapter/user_data/strategies/Utils.py
+++ b/quickadapter/user_data/strategies/Utils.py
@@ -5,7 +5,7 @@ import math
 from enum import IntEnum
 from functools import lru_cache
 from logging import Logger
-from typing import Any, Callable, Literal, Optional, TypeVar, Union
+from typing import Any, Callable, Final, Literal, Optional, TypeVar, Union
 
 import numpy as np
 import optuna
@@ -19,14 +19,14 @@ T = TypeVar("T", pd.Series, float)
 
 
 WeightStrategy = Literal["none", "pivot_threshold"]
-WEIGHT_STRATEGIES: tuple[WeightStrategy, ...] = ("none", "pivot_threshold")
+WEIGHT_STRATEGIES: Final[tuple[WeightStrategy, ...]] = ("none", "pivot_threshold")
 
 NormalizationType = Literal["minmax", "l1", "none"]
-NORMALIZATION_TYPES: tuple[NormalizationType, ...] = ("minmax", "l1", "none")
+NORMALIZATION_TYPES: Final[tuple[NormalizationType, ...]] = ("minmax", "l1", "none")
 
 SmoothingKernel = Literal["gaussian", "kaiser", "triang"]
 SmoothingMethod = Union[SmoothingKernel, Literal["smm", "sma"]]
-SMOOTHING_METHODS: tuple[SmoothingMethod, ...] = (
+SMOOTHING_METHODS: Final[tuple[SmoothingMethod, ...]] = (
     "gaussian",
     "kaiser",
     "triang",
@@ -35,13 +35,13 @@ SMOOTHING_METHODS: tuple[SmoothingMethod, ...] = (
 )
 
 
-DEFAULTS_EXTREMA_SMOOTHING: dict[str, Any] = {
+DEFAULTS_EXTREMA_SMOOTHING: Final[dict[str, Any]] = {
     "method": SMOOTHING_METHODS[0],  # "gaussian"
     "window": 5,
     "beta": 8.0,
 }
 
-DEFAULTS_EXTREMA_WEIGHTING: dict[str, Any] = {
+DEFAULTS_EXTREMA_WEIGHTING: Final[dict[str, Any]] = {
     "normalization": NORMALIZATION_TYPES[0],  # "minmax"
     "gamma": 1.0,
     "strategy": WEIGHT_STRATEGIES[0],  # "none"