diff --git a/lerobot/common/optim/__init__.py b/lerobot/common/optim/__init__.py
index de2c4c99..dd45ed59 100644
--- a/lerobot/common/optim/__init__.py
+++ b/lerobot/common/optim/__init__.py
@@ -12,4 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .optimizers import OptimizerConfig as OptimizerConfig
+from .optimizers import OptimizerConfig
+from .schedulers import LRSchedulerConfig
+
+__all__ = ["OptimizerConfig", "LRSchedulerConfig"]
diff --git a/lerobot/common/policies/config.py b/lerobot/common/policies/config.py
index 167b2ab7..7f4b32ee 100644
--- a/lerobot/common/policies/config.py
+++ b/lerobot/common/policies/config.py
@@ -23,8 +23,7 @@ from huggingface_hub import hf_hub_download
 from huggingface_hub.constants import CONFIG_NAME
 from huggingface_hub.errors import HfHubHTTPError
 
-from lerobot.common.optim.optimizers import OptimizerConfig
-from lerobot.common.optim.schedulers import LRSchedulerConfig
+from lerobot.common.optim import LRSchedulerConfig, OptimizerConfig
 from lerobot.common.utils.hub import HubMixin
 from lerobot.common.utils.utils import auto_select_torch_device, is_amp_available, is_torch_device_available
 from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
diff --git a/lerobot/configs/train.py b/lerobot/configs/train.py
index c9ce6016..6eac0abd 100644
--- a/lerobot/configs/train.py
+++ b/lerobot/configs/train.py
@@ -22,8 +22,7 @@ from huggingface_hub import hf_hub_download
 from huggingface_hub.errors import HfHubHTTPError
 
 from lerobot.common import envs
-from lerobot.common.optim import OptimizerConfig
-from lerobot.common.optim.schedulers import LRSchedulerConfig
+from lerobot.common.optim import LRSchedulerConfig, OptimizerConfig
 from lerobot.common.policies import PreTrainedConfig
 from lerobot.common.utils.hub import HubMixin
 from lerobot.configs import parser