diff --git a/mpc_pybullet_demo/mpcpy/__init__.py b/mpc_pybullet_demo/cvxpy_mpc/__init__.py
similarity index 53%
rename from mpc_pybullet_demo/mpcpy/__init__.py
rename to mpc_pybullet_demo/cvxpy_mpc/__init__.py
index b10be34..52216af 100644
--- a/mpc_pybullet_demo/mpcpy/__init__.py
+++ b/mpc_pybullet_demo/cvxpy_mpc/__init__.py
@@ -1,2 +1,2 @@
-from .cvxpy_mpc import *
+from .cvxpy_mpc import MPC
 from .mpc_config import Params
diff --git a/mpc_pybullet_demo/mpcpy/cvxpy_mpc.py b/mpc_pybullet_demo/cvxpy_mpc/cvxpy_mpc.py
similarity index 72%
rename from mpc_pybullet_demo/mpcpy/cvxpy_mpc.py
rename to mpc_pybullet_demo/cvxpy_mpc/cvxpy_mpc.py
index ca19801..95fff9c 100755
--- a/mpc_pybullet_demo/mpcpy/cvxpy_mpc.py
+++ b/mpc_pybullet_demo/cvxpy_mpc/cvxpy_mpc.py
@@ -2,63 +2,13 @@ import numpy as np
 
 np.seterr(divide="ignore", invalid="ignore")
 
-from scipy.integrate import odeint
-from scipy.interpolate import interp1d
 import cvxpy as opt
 
-from .utils import *
-
 from .mpc_config import Params
 
 P = Params()
 
 
-def get_linear_model_matrices(x_bar, u_bar):
-    """
-    Computes the LTI approximated state space model x' = Ax + Bu + C
-    :param x_bar:
-    :param u_bar:
-    :return:
-    """
-
-    x = x_bar[0]
-    y = x_bar[1]
-    v = x_bar[2]
-    theta = x_bar[3]
-
-    a = u_bar[0]
-    delta = u_bar[1]
-
-    ct = np.cos(theta)
-    st = np.sin(theta)
-    cd = np.cos(delta)
-    td = np.tan(delta)
-
-    A = np.zeros((P.N, P.N))
-    A[0, 2] = ct
-    A[0, 3] = -v * st
-    A[1, 2] = st
-    A[1, 3] = v * ct
-    A[3, 2] = v * td / P.L
-    A_lin = np.eye(P.N) + P.DT * A
-
-    B = np.zeros((P.N, P.M))
-    B[2, 0] = 1
-    B[3, 1] = v / (P.L * cd**2)
-    B_lin = P.DT * B
-
-    f_xu = np.array([v * ct, v * st, a, v * td / P.L]).reshape(P.N, 1)
-    C_lin = (
-        P.DT
-        * (
-            f_xu - np.dot(A, x_bar.reshape(P.N, 1)) - np.dot(B, u_bar.reshape(P.M, 1))
-        ).flatten()
-    )
-
-    # return np.round(A_lin,6), np.round(B_lin,6), np.round(C_lin,6)
-    return A_lin, B_lin, C_lin
-
-
 class MPC:
     def __init__(self, state_cost, final_state_cost, input_cost, input_rate_cost):
         """
@@ -91,20 +41,60 @@ class MPC:
         self.u_bounds = np.array([P.MAX_ACC, P.MAX_STEER])
         self.du_bounds = np.array([P.MAX_D_ACC, P.MAX_D_STEER])
 
-    def optimize_linearized_model(
+    def get_linear_model_matrices(self, x_bar, u_bar):
+        """
+        Computes the LTI approximated state space model x' = Ax + Bu + C
+        :param x_bar:
+        :param u_bar:
+        :return:
+        """
+
+        x = x_bar[0]
+        y = x_bar[1]
+        v = x_bar[2]
+        theta = x_bar[3]
+
+        a = u_bar[0]
+        delta = u_bar[1]
+
+        ct = np.cos(theta)
+        st = np.sin(theta)
+        cd = np.cos(delta)
+        td = np.tan(delta)
+
+        A = np.zeros((self.nx, self.nx))
+        A[0, 2] = ct
+        A[0, 3] = -v * st
+        A[1, 2] = st
+        A[1, 3] = v * ct
+        A[3, 2] = v * td / P.L
+        A_lin = np.eye(self.nx) + self.dt * A
+
+        B = np.zeros((P.N, P.M))
+        B[2, 0] = 1
+        B[3, 1] = v / (P.L * cd**2)
+        B_lin = self.dt * B
+
+        f_xu = np.array([v * ct, v * st, a, v * td / P.L]).reshape(P.N, 1)
+        C_lin = (
+            self.dt
+            * (
+                f_xu
+                - np.dot(A, x_bar.reshape(self.nx, 1))
+                - np.dot(B, u_bar.reshape(self.nu, 1))
+            ).flatten()
+        )
+        return A_lin, B_lin, C_lin
+
+    def step(
         self,
-        A,
-        B,
-        C,
         initial_state,
         target,
+        prev_cmd,
         verbose=False,
     ):
         """
         Optimisation problem defined for the linearised model,
-        :param A:
-        :param B:
-        :param C:
         :param initial_state:
         :param target:
         :param verbose:
@@ -119,6 +109,8 @@ class MPC:
         cost = 0
         constr = []
 
+        A, B, C = self.get_linear_model_matrices(initial_state, prev_cmd)
+
         for k in range(self.control_horizon):
             cost += opt.quad_form(target[:, k] - x[:, k + 1], self.Q)
             cost += opt.quad_form(u[:, k], self.R)
diff --git a/mpc_pybullet_demo/mpcpy/mpc_config.py b/mpc_pybullet_demo/cvxpy_mpc/mpc_config.py
similarity index 100%
rename from mpc_pybullet_demo/mpcpy/mpc_config.py
rename to mpc_pybullet_demo/cvxpy_mpc/mpc_config.py
diff --git a/mpc_pybullet_demo/mpcpy/utils.py b/mpc_pybullet_demo/cvxpy_mpc/utils.py
similarity index 86%
rename from mpc_pybullet_demo/mpcpy/utils.py
rename to mpc_pybullet_demo/cvxpy_mpc/utils.py
index bc5c697..4fa3197 100755
--- a/mpc_pybullet_demo/mpcpy/utils.py
+++ b/mpc_pybullet_demo/cvxpy_mpc/utils.py
@@ -62,19 +62,6 @@ def get_nn_idx(state, path):
     return target_idx
 
 
-def normalize_angle(angle):
-    """
-    Normalize an angle to [-pi, pi]
-    :param angle:
-    :return:
-    """
-    while angle > np.pi:
-        angle -= 2.0 * np.pi
-    while angle < -np.pi:
-        angle += 2.0 * np.pi
-    return angle
-
-
 def get_ref_trajectory(state, path, target_v):
     """
     Reinterpolate the trajectory to get a set N desired target states
@@ -110,5 +97,17 @@ def get_ref_trajectory(state, path, target_v):
     dy = xref[1, :] - state[1]
     xref[0, :] = dx * np.cos(-state[3]) - dy * np.sin(-state[3])  # X
     xref[1, :] = dy * np.cos(-state[3]) + dx * np.sin(-state[3])  # Y
-    xref[3, :] = normalize_angle(path[2, ind] - state[3])  # Theta
+    xref[3, :] = path[2, ind] - state[3]  # Theta
+
+    def fix_angle_reference(angle_ref, angle_init):
+        """
+        This function returns a "smoothened" angle_ref wrt angle_init so there are no jumps.
+        """
+        diff_angle = angle_ref - angle_init
+        diff_angle = np.unwrap(diff_angle)
+        return angle_init + diff_angle
+
+    xref[3, :] = (xref[3, :] + np.pi) % (2.0 * np.pi) - np.pi
+    xref[3, :] = fix_angle_reference(xref[3, :], xref[3, 0])
+
     return xref
diff --git a/mpc_pybullet_demo/mpc_demo_nosim.py b/mpc_pybullet_demo/mpc_demo_nosim.py
index 03978be..e9cad9b 100755
--- a/mpc_pybullet_demo/mpc_demo_nosim.py
+++ b/mpc_pybullet_demo/mpc_demo_nosim.py
@@ -2,14 +2,12 @@
 
 import numpy as np
 import matplotlib.pyplot as plt
-from matplotlib import animation
 from scipy.integrate import odeint
 
-from mpcpy.utils import compute_path_from_wp
-import mpcpy
+from cvxpy_mpc.utils import compute_path_from_wp, get_ref_trajectory
+from cvxpy_mpc import MPC, Params
 
 import sys
-import time
 
 # Robot Starting position
 SIM_START_X = 0.0
@@ -18,7 +16,7 @@ SIM_START_V = 0.0
 SIM_START_H = 0.0
 L = 0.3
 
-params = mpcpy.Params()
+params = Params()
 
 # Params
 VEL = 1.0  # m/s
@@ -44,7 +42,7 @@ class MPCSim:
         R = [10, 10]  # input cost
         P = [10, 10]  # input rate of change cost
 
-        self.mpc = mpcpy.MPC(Q, Qf, R, P)
+        self.mpc = MPC(Q, Qf, R, P)
 
         # Interpolated Path to follow given waypoints
         self.path = compute_path_from_wp(
@@ -111,17 +109,13 @@ class MPCSim:
             # start=time.time()
             # dynamycs w.r.t robot frame
             curr_state = np.array([0, 0, self.state[2], 0])
-            # State Matrices
-            A, B, C = mpcpy.get_linear_model_matrices(curr_state, self.action)
             # Get Reference_traj -> inputs are in worldframe
-            target = mpcpy.get_ref_trajectory(self.state, self.path, VEL)
+            target = get_ref_trajectory(self.state, self.path, VEL)
 
-            x_mpc, u_mpc = self.mpc.optimize_linearized_model(
-                A,
-                B,
-                C,
+            x_mpc, u_mpc = self.mpc.step(
                 curr_state,
                 target,
+                self.action,
                 verbose=False,
             )
             # NOTE: used only for preview purposes
@@ -142,8 +136,8 @@ class MPCSim:
             dxdt = x[2] * np.cos(x[3])
             dydt = x[2] * np.sin(x[3])
             dvdt = u[0]
-            dtheta0dt = x[2] * np.tan(u[1]) / params.L
-            dqdt = [dxdt, dydt, dvdt, dtheta0dt]
+            dthetadt = x[2] * np.tan(u[1]) / params.L
+            dqdt = [dxdt, dydt, dvdt, dthetadt]
             return dqdt
 
         # solve ODE
diff --git a/mpc_pybullet_demo/mpc_demo_pybullet.py b/mpc_pybullet_demo/mpc_demo_pybullet.py
index b3df682..2af6b78 100644
--- a/mpc_pybullet_demo/mpc_demo_pybullet.py
+++ b/mpc_pybullet_demo/mpc_demo_pybullet.py
@@ -1,11 +1,10 @@
 import numpy as np
 import matplotlib.pyplot as plt
-from matplotlib import animation
 
-from mpcpy.utils import compute_path_from_wp
-import mpcpy
+from cvxpy_mpc.utils import compute_path_from_wp, get_ref_trajectory
+from cvxpy_mpc import MPC, Params
 
-params = mpcpy.Params()
+params = Params()
 
 import sys
 import time
@@ -221,7 +220,7 @@ def run_sim():
     R = [10, 10]  # input cost [acc ,steer]
     P = [10, 10]  # input rate of change cost [acc ,steer]
 
-    mpc = mpcpy.MPC(Q, Qf, R, P)
+    mpc = MPC(Q, Qf, R, P)
     x_history = []
     y_history = []
 
@@ -247,7 +246,7 @@ def run_sim():
 
         # Get Reference_traj
         # NOTE: inputs are in world frame
-        target = mpcpy.get_ref_trajectory(state, path, params.TARGET_SPEED)
+        target = get_ref_trajectory(state, path, params.TARGET_SPEED)
 
         # for MPC base link frame is used:
         # so x, y, yaw are 0.0, but speed is the same
@@ -262,16 +261,11 @@ def run_sim():
             ego_state[3] + action[0] * np.tan(action[1]) / params.L * params.DT
         )
 
-        # State Matrices
-        A, B, C = mpcpy.get_linear_model_matrices(ego_state, action)
-
         # optimization loop
         start = time.time()
 
         # MPC step
-        _, u_mpc = mpc.optimize_linearized_model(
-            A, B, C, ego_state, target, verbose=False
-        )
+        _, u_mpc = mpc.step(ego_state, target, action, verbose=False)
         action[0] = u_mpc.value[0, 0]
         action[1] = u_mpc.value[1, 0]