more flexible config

Author: FedeClaudi

proj/model/config.py

@@ -1,54 +1,59 @@
 import numpy as np
 
+# ------------------------------- Model specific params ------------------------------ #
+_cart_params = dict(
+    STATE_SIZE=5,
+    INPUT_SIZE=2,
+    ANGLE_IDX=2,  # state vector index which is angle, used to fit diff in
+    R=np.diag([0.05, 0.05]),  # control cost
+    Q=np.diag([1, 1, 1, 1, 0]),  # state cost | x, y, theta, v, omega
+    Sf=np.diag([0, 0, 0, 0, 0]),  # final state cost
+)
+
+_polar_params = dict(
+    STATE_SIZE=4,
+    INPUT_SIZE=2,
+    ANGLE_IDX=2,  # state vector index which is angle, used to fit diff in
+    R=np.diag([0.05, 0.05]),  # control cost
+    Q=np.diag([2.5, 2.5, 0, 0]),  # state cost | r, omega, v, omega
+    Sf=np.diag([2.5, 2.5, 0, 0]),  # final state cost
+)
+
+# -------------------------------- Mouse types ------------------------------- #
+
+_easy_mouse = dict(
+    L=1.5,  # half body width | cm
+    R=1,  # radius of wheels | cm
+    d=0.1,  # distance between axel and CoM | cm
+    length=3,  # cm
+    m=round(20 / 9.81, 2),  # mass | g
+    m_w=round(2 / 9.81, 2),  # mass of wheels/legs |g
+    mouse_type="easy",
+)
+
+_realistic_mouse = dict(
+    L=2,  # half body width | cm
+    R=2,  # radius of wheels | cm
+    d=3,  # distance between axel and CoM | cm
+    length=8.6,  # cm
+    m=round(25 / 9.81, 2),  # mass | g
+    m_w=round(0.6 / 9.81, 2),  # mass of wheels/legs |g
+    mouse_type="realistic",
+)
+
 
 class Config:
+    # ----------------------------- Simulation params ---------------------------- #
     SIMULATION_NAME = ""
 
     USE_FAST = True  # if true use cumba's methods
     SPAWN_TYPE = "trajectory"
     LIVE_PLOT = True
 
-    # ----------------------------- Simulation params ---------------------------- #
-    dt = 0.005
-
-    # -------------------------------- Cost params ------------------------------- #
-    STATE_SIZE = 5
-    INPUT_SIZE = 2
-    ANGLE_IDX = 2  # state vector index which is angle, used to fit diff in
-
-    R = np.diag([0.05, 0.05])  # control cost
-    Q = np.diag([1, 1, 1, 1, 0])  # state cost | x, y, theta, v, omega
-    Sf = np.diag([0, 0, 0, 0, 0])  # final state cost
-
-    # STATE_SIZE = 4
-    # INPUT_SIZE = 2
-
-    # R = np.diag([0.05, 0.05])  # control cost
-    # Q = np.diag([2.5, 2.5, 0, 0])  # state cost | r, omega, v, omega
-    # Sf = np.diag([2.5, 2.5, 0, 0])  # final state cost
-
-    # ------------------------------- Mouse params ------------------------------- #
-    # ? works
-    mouse = dict(
-        mouse_type="working",
-        L=1.5,  # half body width | cm
-        R=1,  # radius of wheels | cm
-        d=0.1,  # distance between axel and CoM | cm
-        length=3,  # cm
-        m=round(20 / 9.81, 2),  # mass | g
-        m_w=round(2 / 9.81, 2),  # mass of wheels/legs |g
-    )
+    mouse_type = "easy"
+    model_type = "cart"
 
-    # ? more realistic
-    # mouse = dict(
-    # mouse_type = 'realistic',
-    #     L=2,  # half body width | cm
-    #     R=2,  # radius of wheels | cm
-    #     d=3,  # distance between axel and CoM | cm
-    #     length=8.6,  # cm
-    #     m=round(25 / 9.81, 2),  # mass | g
-    #     m_w=round(0.6 / 9.81, 2),  # mass of wheels/legs |g
-    # )
+    dt = 0.005
 
     # ------------------------------ Goal trajectory ----------------------------- #
 
@@ -84,6 +89,25 @@ class Config:
         threshold=1e-6,
     )
 
+    def __init__(self,):
+        # get mouse params
+        self.mouse = (
+            _easy_mouse if self.mouse_type == "easy" else _realistic_mouse
+        )
+
+        # set model params
+        if self.model_type == "cart":
+            params = _cart_params
+        else:
+            params = _polar_params
+
+        self.STATE_SIZE = params["STATE_SIZE"]
+        self.INPUT_SIZE = params["INPUT_SIZE"]
+        self.ANGLE_IDX = params["ANGLE_IDX"]
+        self.R = params["R"]
+        self.Q = params["Q"]
+        self.Sf = params["Sf"]
+
     def config_dict(self):
         return dict(
             dt=self.dt,