Merge pull request #16 from ExcitingSystems/enum

Add enums

Author: hvater

README.md

@@ -9,14 +9,14 @@ A basic routine is as simple as:
 ```py
 import jax.numpy as jnp
 import exciting_environments as excenvs
+from exciting_environments import EnvironmentRegistry
 from exciting_environments.utils import MinMaxNormalization
 
-env = excenvs.make(
-    "Pendulum-v0",
+env = EnvironmentRegistry.PENDULUM.make(
     batch_size=5,
     action_normalizations={"torque": MinMaxNormalization(min=-15,max=15)},
     tau=2e-2
-) 
+)
 obs, state = env.vmap_reset()
 
 actions = jnp.linspace(start=-1, stop=1, num=1000)[None, :, None]
@@ -45,11 +45,11 @@ alternatively, simulate full trajectories:
 ```py
 import jax.numpy as jnp
 import exciting_environments as excenvs
+from exciting_environments import EnvironmentRegistry
 from exciting_environments.utils import MinMaxNormalization
 import diffrax
 
-env = excenvs.make(
-    "Pendulum-v0",
+env = EnvironmentRegistry.PENDULUM.make(
     solver=diffrax.Tsit5(),
     batch_size=5,
     action_normalizations={"torque": MinMaxNormalization(min=-15,max=15)},
@@ -76,4 +76,4 @@ which produces $5$ identical trajectories in parallel as well:
 ![](https://github.com/ExcitingSystems/exciting-environments/blob/main/fig/excenvs_pendulum_simulation_example_advanced.png?raw=true)
 
 Note that in this case the Tsit5 ODE solver instead of the default explicit Euler is used.
-All solvers used here are from the diffrax library (https://docs.kidger.site/diffrax/).
+All solvers used here are from the diffrax library (https://docs.kidger.site/diffrax/).

examples/pmsm_example.ipynb

@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -14,6 +14,7 @@
     "from typing import Callable\n",
     "import os\n",
     "from exciting_environments import PMSM\n",
+    "from exciting_environments.pmsm import MotorVariant\n",
     "import matplotlib.pyplot as plt\n",
     "import jax_dataclasses as jdc\n",
     "from exciting_environments.utils import MinMaxNormalization\n",
@@ -22,14 +23,14 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "BATCH_SIZE=5\n",
     "new_motor_env = PMSM(\n",
     "    saturated=True,\n",
-    "    LUT_motor_name=\"BRUSA\",\n",
+    "    motor_variant=MotorVariant.BRUSA,\n",
     "    batch_size=BATCH_SIZE,\n",
     "    control_state=[\"torque\"]\n",
     "    )"
@@ -48,7 +49,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -84,7 +85,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -138,7 +139,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3",
+   "display_name": "venv_dev",
    "language": "python",
    "name": "python3"
   },
@@ -152,7 +153,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.3"
+   "version": "3.11.9"
   },
   "orig_nbformat": 4
  },

exciting_environments/__init__.py

@@ -5,6 +5,6 @@
 from .pendulum import Pendulum
 from .fluid_tank import FluidTank
 from .acrobot import Acrobot
-from .registration import make
+from .registration import EnvironmentRegistry
 from .gym_wrapper import GymWrapper
 from .mujoco_wrapper import MujucoWrapper

exciting_environments/gym_wrapper.py

@@ -6,8 +6,9 @@
 from functools import partial
 import chex
 from abc import ABC
-from exciting_environments import spaces
-from exciting_environments.registration import make
+from exciting_environments import spaces, EnvironmentRegistry
+
+# from exciting_environments.registration import make
 
 
 class GymWrapper(ABC):
@@ -58,9 +59,9 @@ def __init__(
             self.generate_terminated = self.env.generate_terminated
 
     @classmethod
-    def from_name(cls, env_id: str, **env_kwargs):
-        """Creates GymWrapper with environment based on passed env_id."""
-        env = make(env_id, **env_kwargs)
+    def from_env(cls, env_type: EnvironmentRegistry, **env_kwargs):
+        """Creates GymWrapper with environment from EnvironmentRegistry."""
+        env = env_type.make(**env_kwargs)
         return cls(env)
 
     def step(self, action):

exciting_environments/pmsm/__init__.py

@@ -1,2 +1,2 @@
-from .motor_parameters import default_params
+from .motor_parameters import MotorVariant
 from .pmsm_env import PMSM  # , PMSM_Physical

exciting_environments/pmsm/motor_parameters.py

@@ -10,6 +10,8 @@
 from exciting_environments.utils import MinMaxNormalization
 from copy import deepcopy
 
+from enum import Enum
+
 
 @jdc.pytree_dataclass
 class PhysicalNormalizations:
@@ -147,21 +149,15 @@ def default_soft_constraints(self, state, action_norm, env_properties):
 )
 
 
-def default_params(name):
-    """
-    Returns default parameters for specified motor configurations.
-
-    Args:
-        name (str): Name of the motor ("BRUSA" or "SEW").
-
-    Returns:
-        MotorConfig: Configuration containing physical constraints, action constraints, static parameters, and LUT data.
-    """
-    if name is None:
-        return deepcopy(DEFAULT)
-    elif name == "BRUSA":
-        return deepcopy(BRUSA)
-    elif name == "SEW":
-        return deepcopy(SEW)
-    else:
-        raise ValueError(f"Motor name {name} is not known.")
+class MotorVariant(Enum):
+    DEFAULT = "DEFAULT"
+    BRUSA = "BRUSA"
+    SEW = "SEW"
+
+    def get_params(self):
+        if self is MotorVariant.BRUSA:
+            return deepcopy(BRUSA)
+        elif self is MotorVariant.SEW:
+            return deepcopy(SEW)
+        else:
+            return deepcopy(DEFAULT)

exciting_environments/pmsm/pmsm_env.py

@@ -14,7 +14,7 @@
 from copy import deepcopy
 
 from exciting_environments import CoreEnvironment
-from exciting_environments.pmsm import default_params
+from exciting_environments.pmsm import MotorVariant
 
 
 # only for alpha/beta -> abc
@@ -117,7 +117,7 @@ def __init__(
         self,
         batch_size: int = 8,
         saturated=False,
-        LUT_motor_name: str = None,
+        motor_variant: MotorVariant = MotorVariant.DEFAULT,
         physical_normalizations: dict = None,
         action_normalizations: dict = None,
         soft_constraints: Callable = None,
@@ -130,7 +130,7 @@ def __init__(
         Args:
             batch_size (int): Number of parallel environment simulations. Default: 8
             saturated (bool): Permanent magnet flux linkages and inductances are taken from LUT_motor_name specific LUTs. Default: False
-            LUT_motor_name (str): Sets physical_normalizations, action_normalizations, soft_constraints and static_params to default values for the passed motor name and stores associated LUTs for the possible saturated case. Needed if saturated==True.
+            motor_variant (MotorVariant): Sets physical_normalizations, action_normalizations, soft_constraints and static_params to default values for the passed motor variant and stores associated LUTs for the possible saturated case. Needed if saturated==True.
             physical_normalizations (dict): min-max normalization values of the physical state of the environment.
                 u_d_buffer (MinMaxNormalization): Direct share of the delayed action due to system deadtime. Default: min=-2 * 400 / 3, max=2 * 400 / 3
                 u_q_buffer (MinMaxNormalization): Quadrature share of the delayed action due to system deadtime. Default: min=-2 * 400 / 3, max=2 * 400 / 3
@@ -161,8 +161,8 @@ def __init__(
         self.tau = tau
         self._solver = solver
 
-        if LUT_motor_name is not None:
-            motor_params = deepcopy(default_params(LUT_motor_name))
+        if motor_variant != MotorVariant.DEFAULT:
+            motor_params = motor_variant.get_params()
             default_physical_normalizations = motor_params.physical_normalizations.__dict__
             default_action_normalizations = motor_params.action_normalizations.__dict__
             default_static_params = motor_params.static_params.__dict__
@@ -187,7 +187,10 @@ def __init__(
 
         else:
             if saturated:
-                raise Exception("LUT_motor_name is needed to load LUTs.")
+                raise ValueError(
+                    f"MotorVariant '{motor_variant.value}' is not allowed for saturated LUTs. "
+                    "Use a specific motor variant. DEFAULT is only valid for saturated=False."
+                )
 
             saturated_quants = [
                 "L_dd",
@@ -198,7 +201,7 @@ def __init__(
                 "Psi_q",
             ]
 
-            motor_params = deepcopy(default_params(LUT_motor_name))
+            motor_params = motor_variant.get_params()
             default_physical_normalizations = motor_params.physical_normalizations.__dict__
             default_action_normalizations = motor_params.action_normalizations.__dict__
             default_static_params = motor_params.static_params.__dict__

exciting_environments/registration.py

@@ -6,29 +6,27 @@
     PMSM,
     Acrobot,
 )
-
-
-def make(env_id: str, **env_kwargs):
-    if env_id == "CartPole-v0":
-        env = CartPole(**env_kwargs)
-
-    elif env_id == "MassSpringDamper-v0":
-        env = MassSpringDamper(**env_kwargs)
-
-    elif env_id == "Pendulum-v0":
-        env = Pendulum(**env_kwargs)
-
-    elif env_id == "FluidTank-v0":
-        env = FluidTank(**env_kwargs)
-
-    elif env_id == "PMSM-v0":
-        env = PMSM(**env_kwargs)
-
-    elif env_id == "Acrobot-v0":
-        env = Acrobot(**env_kwargs)
-
-    else:
-        print(f"No existing environments got env_id ={env_id}")
-        env = None
-
-    return env
+from enum import Enum
+
+
+class EnvironmentRegistry(Enum):
+    CART_POLE = "CartPole-v0"
+    MASS_SPRING_DAMPER = "MassSpringDamper-v0"
+    PENDULUM = "Pendulum-v0"
+    FLUID_TANK = "FluidTank-v0"
+    PMSM = "PMSM-v0"
+    ACROBOT = "Acrobot-v0"
+
+    def make(self, **env_kwargs):
+        env_map = {
+            EnvironmentRegistry.CART_POLE: CartPole,
+            EnvironmentRegistry.MASS_SPRING_DAMPER: MassSpringDamper,
+            EnvironmentRegistry.PENDULUM: Pendulum,
+            EnvironmentRegistry.FLUID_TANK: FluidTank,
+            EnvironmentRegistry.PMSM: PMSM,
+            EnvironmentRegistry.ACROBOT: Acrobot,
+        }
+        cls = env_map.get(self)
+        if cls is None:
+            raise ValueError(f"Unknown environment: {self}")
+        return cls(**env_kwargs)

tests/envs/acrobot/test_acrobot.py

@@ -4,11 +4,13 @@
 import jax.numpy as jnp
 import numpy as np
 import diffrax
+from exciting_environments import EnvironmentRegistry
 from exciting_environments.utils import MinMaxNormalization, load_sim_properties_from_json
 from pathlib import Path
 import pickle
 import os
 
+
 jax.config.update("jax_enable_x64", True)
 
 
@@ -33,7 +35,7 @@ def test_default_initialization():
         "omega_1": MinMaxNormalization(min=-10, max=10),
         "omega_2": MinMaxNormalization(min=-10, max=10),
     }
-    env = excenvs.make("Acrobot-v0", batch_size=batch_size)
+    env = EnvironmentRegistry.ACROBOT.make(batch_size=batch_size)
     for key, value in params.items():
         env_value = getattr(env.env_properties.static_params, key)
         if isinstance(value, jnp.ndarray) or isinstance(env_value, jnp.ndarray):
@@ -101,8 +103,7 @@ def test_custom_initialization():
         "omega_1": MinMaxNormalization(min=-55, max=10),
         "omega_2": MinMaxNormalization(min=-10, max=30),
     }
-    env = excenvs.make(
-        "Acrobot-v0",
+    env = EnvironmentRegistry.ACROBOT.make(
         batch_size=batch_size,
         static_params=params,
         physical_normalizations=physical_normalizations,
@@ -159,8 +160,7 @@ def test_step_results():
     loaded_params, loaded_action_normalizations, loaded_physical_normalizations, loaded_tau = (
         load_sim_properties_from_json(file_path)
     )
-    env = excenvs.make(
-        "Acrobot-v0",
+    env = EnvironmentRegistry.ACROBOT.make(
         tau=loaded_tau,
         solver=diffrax.Euler(),
         static_params=loaded_params,