adding wheels kinematics

Author: FedeClaudi

proj/model/config.py

@@ -52,7 +52,7 @@ class Config:
 
     USE_FAST = True  # if true use cumba's methods
     SPAWN_TYPE = "trajectory"
-    LIVE_PLOT = False
+    LIVE_PLOT = True
 
     mouse_type = "realistic"
     model_type = "cart"
@@ -76,7 +76,7 @@ class Config:
 
     # ------------------------------ Planning params ----------------------------- #
     planning = dict(  # params used to compute goal states to be used for control
-        prediction_length=20,
+        prediction_length=80,
         n_ahead=5,  # start prediction states from N steps ahead
     )
 

proj/model/model.py

@@ -55,7 +55,8 @@ class Model(Config):
 
     _control = namedtuple("control", "tau_r, tau_l")
     _state = namedtuple("state", "x, y, theta, v, omega")
-    _dxdt = namedtuple("dxdt", "x_dot, y_dt, theta_dot, v_dot, omega_dot")
+    _dxdt = namedtuple("dxdt", "x_dot, y_dot, theta_dot, v_dot, omega_dot")
+    _wheel_state = namedtuple("wheel_state", "nudot_right, nudot_left")
 
     def __init__(self, startup=True):
         Config.__init__(self)
@@ -71,6 +72,8 @@ def __init__(self, startup=True):
                 self.calc_dqdt = fast_dqdt
                 self.calc_model_jacobian_state = fast_model_jacobian_state
                 self.calc_model_jacobian_input = fast_model_jacobian_input
+
+            self.get_wheels_dynamics()
             self.reset()
 
     def reset(self):
@@ -85,10 +88,12 @@ def reset(self):
             r=[],
             gamma=[],
             trajectory_idx=[],
+            nudot_left=[],  # acceleration of left wheel
+            nudot_right=[],  # acceleration of right wheel
         )
 
     def _append_history(self):
-        for ntuple in [self.curr_x, self.curr_control]:
+        for ntuple in [self.curr_x, self.curr_control, self.curr_wheel_state]:
             for k, v in ntuple._asdict().items():
                 self.history[k].append(v)
 
@@ -251,6 +256,43 @@ def get_jacobians(self):
             args, self.model_jacobian_input, modules="numpy"
         )
 
+    def get_wheels_dynamics(self):
+        (
+            x,
+            y,
+            theta,
+            L,
+            R,
+            m,
+            m_w,
+            d,
+            tau_l,
+            tau_r,
+            v,
+            omega,
+        ) = self.variables.values()
+
+        x_dot, y_dot, theta_dot = symbols("xdot, ydot, thetadot")
+        nu_l_dot, nu_r_dot = symbols("nudot_L, nudot_R")
+
+        vels = Matrix([x_dot, y_dot, theta_dot])
+        K = Matrix(
+            [
+                [R / 2 * cos(theta), R / 2 * cos(theta)],
+                [R / 2 * sin(theta), R / 2 * sin(theta)],
+                [R / (2 * L), R / (2 * L)],
+            ]
+        )
+
+        # In the model you can use the wheels accelerations
+        # to get the x,y,theta velocity.
+        # Here we do the inverse, given x,y,theta velocities
+        # we get the wheel's accelerations
+
+        nu = K.pinv() * vels
+        args = [L, R, theta, x_dot, y_dot, theta_dot]
+        self.calc_wheels_accels = lambdify(args, nu, modules="numpy")
+
     def step(self, u):
         u = self._control(*np.array(u))
         self.curr_x = self._state(*self.curr_x)
@@ -268,6 +310,17 @@ def step(self, u):
         next_x = np.array(self.curr_x) + dxdt * self.dt
         self.curr_x = self._state(*next_x)
 
+        # Compute wheel accelerations
+        w = self.calc_wheels_accels(
+            variables["L"],
+            variables["R"],
+            variables["theta"],
+            self.curr_dxdt.x_dot,
+            self.curr_dxdt.y_dot,
+            self.curr_dxdt.theta_dot,
+        )
+        self.curr_wheel_state = self._wheel_state(*w.ravel())
+
         # Update history
         self.curr_control = u
         self._append_history()

proj/plotting/results.py

@@ -135,21 +135,21 @@ def _plot_cost(cost_history, ax=None):
 
 
 def _plot_integrals(history, dt, tax=None, aax=None):
-    R, L = np.cumsum(history["tau_r"]), np.cumsum(history["tau_l"])
+    R, L = history["nudot_right"], history["nudot_left"]
 
     plot_line_outlined(
         tax,
         R,
         color=desaturate_color(colors["tau_r"]),
-        label="R_wheel_speed",
+        label="R_wheel_accel",
         lw=2,
         solid_capstyle="round",
     )
     plot_line_outlined(
         tax,
         L,
         color=desaturate_color(colors["tau_l"]),
-        label="L_wheel_speed",
+        label="L_wheel_accel",
         lw=2,
         solid_capstyle="round",
     )

workspace.py

@@ -1,21 +1,43 @@
 # %%
-from proj.utils.misc import load_results_from_folder
+from proj import Model
+from sympy import Matrix, symbols, sin, cos
 
-fld = "/Users/federicoclaudi/Dropbox (UCL - SWC)/Rotation_vte/Locomotion/control/tracking_200923_141709_1491_good_example/results"
+model = Model()
 
 # %%
-config, trajectory, history, cost_history = load_results_from_folder(fld)
+(
+    x,
+    y,
+    theta,
+    L,
+    R,
+    m,
+    m_w,
+    d,
+    tau_l,
+    tau_r,
+    v,
+    omega,
+) = model.variables.values()
 
 # %%
-import matplotlib.pyplot as plt
+x_dot, y_dot, theta_dot = symbols("xdot, ydot, thetadot")
 
-plt.plot(cost_history["x"])
+nu_l_dot, nu_r_dot = symbols("nudot_L, nudot_R")
 
-# %%
-plt.plot(history["tau_r"])
-plt.plot(history["tau_l"])
 
-# %%
-plt.plot(history["v"])
+vels = Matrix([x_dot, y_dot, theta_dot])
+
+K = Matrix(
+    [
+        [R / 2 * cos(theta), R / 2 * cos(theta)],
+        [R / 2 * sin(theta), R / 2 * sin(theta)],
+        [R / (2 * L), R / (2 * L)],
+    ]
+)
+
+nu = Matrix([nu_l_dot, nu_r_dot])
 
+nu = K.pinv() * vels
+nu
 # %%