diff --git a/matlab/LQR_infHor.m b/matlab/LQR_infHor.m
index 6604df34f29a7895993116cd5d8cc1f324f0158a..386c0afa0abc0f1b668be47f10343748fff4b4ae 100644
--- a/matlab/LQR_infHor.m
+++ b/matlab/LQR_infHor.m
@@ -1,7 +1,7 @@
-%% Point-mass LQR with infinite horizon
+%% Point-mass discrete LQR with infinite horizon
 %% 
-%% Copyright (c) 2023 Idiap Research Institute <https://www.idiap.ch/>
-%% Written by Sylvain Calinon <https://calinon.ch>
+%% Copyright (c) 2024 Idiap Research Institute <https://www.idiap.ch/>
+%% Written by Ante Marić <ante.maric@idiap.ch> and Sylvain Calinon <https://calinon.ch>
 %% 
 %% This file is part of RCFS <https://robotics-codes-from-scratch.github.io/>
 %% License: GPL-3.0-only
@@ -21,12 +21,15 @@ param.rfactor = 4E-2; %Control cost in LQR
 %Control cost matrix
 R = eye(param.nbVarPos) * param.rfactor;
 
-%Target and desired covariance
+%Target and desired covariance (for precision matrix on static features)
 param.Mu = [randn(param.nbVarPos,1); zeros(param.nbVarPos*(param.nbDeriv-1),1)];
 
 [Ar,~] = qr(randn(param.nbVarPos));
 xCov = Ar * diag(rand(param.nbVarPos,1)) * Ar' * 1E-1;
 
+%Target and desired covariance (for precision matrix on static and dynamic features)
+%param.Mu = [randn(param.nbVarPos,1); zeros(param.nbVarPos*(param.nbDeriv-1),1)];
+%xCov = diag([rand(2,1); rand(2,1) / param.dt]) * 1E-1;
 
 %% Discrete dynamical System settings 
 %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
@@ -44,7 +47,12 @@ B = kron(B1d, eye(param.nbVarPos)); %Discrete nD
 
 %% discrete LQR with infinite horizon
 %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
-Q = blkdiag(inv(xCov), zeros(param.nbVarPos*(param.nbDeriv-1))); %Precision matrix
+% Precision matrix on static features
+Q = blkdiag(inv(xCov), zeros(param.nbVarPos*(param.nbDeriv-1)));
+
+% Precision matrix on static and dynamic features
+% Q = inv(xCov);
+
 P = solveAlgebraicRiccati_eig_discrete(A, B*(R\B'), (Q+Q')/2);
 L = (B' * P * B + R) \ B' * P * A; %Feedback gain (discrete version)
 
@@ -59,7 +67,7 @@ for n=1:param.nbRepros
     x = [ones(param.nbVarPos,1)+randn(param.nbVarPos,1)*5E-1; zeros(param.nbVarPos*(param.nbDeriv-1),1)];
     for t=1:param.nbData		
         r(n).Data(:,t) = x; 
-        u = L * (param.Mu - x); %Compute acceleration (with only feedback terms)
+        u = L * (param.Mu - x); %Compute control input (with only feedback terms)
         x = A * x + B * u;
     end
 end
@@ -110,7 +118,7 @@ function X = solveAlgebraicRiccati_eig_discrete(A, G, Q)
     n = size(A,1);
 
     %Symplectic matrix (see https://en.wikipedia.org/wiki/Algebraic_Riccati_equation)
-    %Z = [A+B*(R\B')/A'*Q, -B*(R\B')/A'; -A'\Q, A'^-1]; 
+    %Z = [A+B*(R\B')/A'*Q, -B*(R\B')/A'; -A'\Q, A'^-1];
     Z = [A+G/A'*Q, -G/A'; -A'\Q, inv(A')];
 
     %Since Z is symplectic, if it does not have any eigenvalues on the unit circle, 
diff --git a/python/LQR_infHor.py b/python/LQR_infHor.py
index 4ea88c1196d2ec66dfa2095c1a0ca9c950530e68..b125afa05c13777cd35c58d335b2fd73ec6114d5 100644
--- a/python/LQR_infHor.py
+++ b/python/LQR_infHor.py
@@ -1,8 +1,8 @@
 '''
-Point-mass LQR with infinite horizon
+Point-mass discrete LQR with infinite horizon
 
-Copyright (c) 2023 Idiap Research Institute <https://www.idiap.ch/>
-Written by Jérémy Maceiras <jeremy.maceiras@idiap.ch> and
+Copyright (c) 2024 Idiap Research Institute <https://www.idiap.ch/>
+Written by Ante Marić <ante.maric@idiap.ch>, Jérémy Maceiras <jeremy.maceiras@idiap.ch> and
 Sylvain Calinon <https://calinon.ch>
 
 This file is part of RCFS <https://robotics-codes-from-scratch.github.io/>
@@ -11,83 +11,119 @@ License: GPL-3.0-only
 
 import numpy as np
 import copy
-from scipy.linalg import solve_discrete_are as solve_algebraic_riccati_discrete
-from scipy.stats import multivariate_normal
 from math import factorial
 import matplotlib.pyplot as plt
-
-# Plot a 2D Gaussians
-def plot_gaussian(mu, sigma):
-    w, h = 100, 100
-
-    std = [np.sqrt(sigma[0, 0]), np.sqrt(sigma[1, 1])]
-    x = np.linspace(mu[0] - 3 * std[0], mu[0] + 3 * std[0], w)
-    y = np.linspace(mu[1] - 3 * std[1], mu[1] + 3 * std[1], h)
-
-    x, y = np.meshgrid(x, y)
-
-    x_ = x.flatten()
-    y_ = y.flatten()
-    xy = np.vstack((x_, y_)).T
-
-    normal_rv = multivariate_normal(mu, sigma)
-    z = normal_rv.pdf(xy)
-    z = z.reshape(w, h, order='F')
-
-    plt.contourf(x, y, z.T,levels=1,colors=["white","red"],alpha=.5)
+import matplotlib.patches as patches
+
+def plot_gaussian(Mu, Sigma, color, valAlpha=1):
+    nbDrawingSeg = 100
+    darkcolor = color * 0.7
+    t = np.linspace(-np.pi, np.pi, nbDrawingSeg)
+
+    D, V = np.linalg.eig(Sigma)
+    R = np.real(V @ np.sqrt(np.diag(D)))
+    X = R @ np.array([np.cos(t), np.sin(t)]) + Mu.reshape(-1, 1)
+
+    if valAlpha < 1:  # Plot with alpha transparency
+        patch = patches.Polygon(X.T, closed=True, facecolor=color, edgecolor=color, alpha=valAlpha, zorder=0)
+        plt.gca().add_patch(patch)
+        plt.plot(Mu[0], Mu[1], '.', markersize=10, color=darkcolor, zorder=1)
+    else:  # Plot without transparency
+        patch = patches.Polygon(X.T, closed=True, facecolor=color, edgecolor=darkcolor, zorder=0)
+        plt.gca().add_patch(patch)
+        plt.plot(Mu[0], Mu[1], '.', markersize=10, color=darkcolor, zorder=1)
+
+    plt.axis('equal')
+
+def solve_algebraic_riccati_eig_discrete(A, G, Q):
+    n = A.shape[0]
+
+    # Symplectic matrix (see https://en.wikipedia.org/wiki/Algebraic_Riccati_equation)
+    Z = np.block([
+        [A + G @ np.linalg.inv(A).T @ Q, -G @ np.linalg.inv(A).T],
+        [-np.linalg.inv(A).T @ Q, np.linalg.inv(A).T]
+    ])
+
+    # Since Z is symplectic, if it does not have any eigenvalues on the unit circle,
+    # then exactly half of its eigenvalues are inside the unit circle.
+    eigvals, eigvecs = np.linalg.eig(Z)
+    U = []
+    for j in range(2 * n):
+        if np.abs(eigvals[j]) < 1:  # inside unit circle
+            U.append(eigvecs[:, j])
+
+    U = np.array(U).T
+    X = np.real(U[n:, :] @ np.linalg.inv(U[:n, :]))
+    return X
 
 # Parameters
 # ===============================
 
-param = lambda: None # Lazy way to define an empty class in python
-param.nbData = 200 # Number of datapoints
-param.nbRepros = 4 #Number of reproductions
-param.nbVarPos = 2 #Dimension of position data (here: x1,x2)
-param.nbDeriv = 2 #Number of static & dynamic features (D=2 for [x,dx])
-param.nbVar = param.nbVarPos * param.nbDeriv #Dimension of state vector in the tangent space
-param.dt = 1e-2 #Time step duration
-param.rfactor = 4e-2	#Control cost in LQR 
+param = lambda: None  # Lazy way to define an empty class in python
+param.nbData = 200  # Number of datapoints
+param.nbRepros = 4  # Number of reproductions
+param.nbVarPos = 2  # Dimension of position data (here: x1,x2)
+param.nbDeriv = 2  # Number of static & dynamic features (D=2 for [x,dx])
+param.nbVar = param.nbVarPos * param.nbDeriv  # Dimension of state vector in the tangent space
+param.dt = 1e-2  # Time step duration
+param.rfactor = 4e-2  # Control cost in LQR
 
 # Control cost matrix
 R = np.identity(param.nbVarPos) * param.rfactor
 
-# Target and desired covariance
-param.Mu = np.hstack(( np.random.uniform(size=param.nbVarPos) , np.zeros(param.nbVarPos) ))
-Ar,_ = np.linalg.qr(np.random.uniform(size=(param.nbVarPos,param.nbVarPos)))
+# Target and desired covariance (for precision matrix on static features)
+param.Mu = np.hstack((np.random.uniform(size=param.nbVarPos), np.zeros(param.nbVarPos)))
+Ar, _ = np.linalg.qr(np.random.uniform(size=(param.nbVarPos, param.nbVarPos)))
 xCov = Ar @ np.diag(np.random.uniform(size=param.nbVarPos)) @ Ar.T * 1e-1
 
+## Target and desired covariance (for precision matrix on static and dynamic features)
+# param.Mu = np.hstack((np.random.randn(param.nbVarPos), np.zeros(param.nbVarPos * (param.nbDeriv - 1))))
+# xCov = np.diag(np.hstack((np.random.rand(2), np.random.rand(2) / param.dt))) * 1e-1
+
 # Discrete dynamical System settings 
 # ===================================
 
-A1d = np.zeros((param.nbDeriv,param.nbDeriv))
-B1d = np.zeros((param.nbDeriv,1))
+A1d = np.zeros((param.nbDeriv, param.nbDeriv))
+B1d = np.zeros((param.nbDeriv, 1))
 
 for i in range(param.nbDeriv):
-    A1d += np.diag( np.ones(param.nbDeriv-i) ,i ) * param.dt**i * 1/factorial(i)
-    B1d[param.nbDeriv-i-1] = param.dt**(i+1) * 1/factorial(i+1)
+    A1d += np.diag(np.ones(param.nbDeriv - i), i) * param.dt**i * 1 / factorial(i)
+    B1d[param.nbDeriv - i - 1] = param.dt**(i + 1) * 1 / factorial(i + 1)
 
-A = np.kron(A1d,np.identity(param.nbVarPos))
-B = np.kron(B1d,np.identity(param.nbVarPos))
+A = np.kron(A1d, np.identity(param.nbVarPos))
+B = np.kron(B1d, np.identity(param.nbVarPos))
 
 # Discrete LQR with infinite horizon
 # ===================================
 
-Q = np.zeros((param.nbVar,param.nbVar))
-Q[:param.nbVarPos,:param.nbVarPos] = np.linalg.inv(xCov) # Precision matrix
-P = solve_algebraic_riccati_discrete(A,B,Q,R)
-L = np.linalg.inv(B.T @ P @ B + R) @ B.T @ P @ A # Feedback gain (discrete version)
+# Precision matrix on static features
+Q = np.zeros((param.nbVar, param.nbVar))
+Q[:param.nbVarPos, :param.nbVarPos] = np.linalg.inv(xCov)  # Precision matrix
+
+## Precision matrix on static and dynamic features
+# Q = np.linalg.inv(xCov)
+
+P = solve_algebraic_riccati_eig_discrete(A, B @ np.linalg.inv(R) @ B.T, (Q + Q.T) / 2)
+L = np.linalg.inv(B.T @ P @ B + R) @ B.T @ P @ A  # Feedback gain (discrete version)
+
+# Test ratio between kp and kv
+if param.nbDeriv > 1:
+    kp = np.linalg.eigvals(L[:, :param.nbVarPos])
+    kv = np.linalg.eigvals(L[:, param.nbVarPos:])
+    ratio = kv / (2 * np.sqrt(kp))  # Corresponds to ideal damping ratio
+    print("Ratio:", ratio)
 
-reproducitons = []
+reproductions = []
 
 for i in range(param.nbRepros):
     xt = np.zeros(param.nbVar)
-    xt[:param.nbVarPos] = 1+np.random.uniform(param.nbVarPos)*2 
+    xt[:param.nbVarPos] = 1 + np.random.uniform(param.nbVarPos) * 2
     xs = [copy.deepcopy(xt)]
     for t in range(param.nbData):
         u = L @ (param.Mu - xt)
         xt = A @ xt + B @ u
         xs += [copy.deepcopy(xt)]
-    reproducitons += [ np.asarray(xs) ]
+    reproductions += [np.asarray(xs)]
 
 # Plots
 # ======
@@ -95,20 +131,20 @@ for i in range(param.nbRepros):
 plt.figure()
 plt.title("Position")
 
-for r in reproducitons:
-    plt.plot(r[:,0],r[:,1],c="black")
-plot_gaussian(param.Mu[:param.nbVarPos],xCov)
+for r in reproductions:
+    plt.plot(r[:, 0], r[:, 1], c="black")
+plot_gaussian(param.Mu[:param.nbVarPos], xCov[:2, :2], color=np.array([.8, 0, 0]), valAlpha=0.3)
 plt.axis("off")
 plt.gca().set_aspect('equal', adjustable='box')
 
-fig,axs = plt.subplots(5)
+fig, axs = plt.subplots(5)
 
-for r in reproducitons:
-    axs[0].plot(r[:,0],c="black",alpha=.4,linestyle="dashed")
-    axs[1].plot(r[:,1],c="black",alpha=.4,linestyle="dashed")    
-    axs[2].plot(r[:,2],c="black",alpha=.4,linestyle="dashed")
-    axs[3].plot(r[:,3],c="black",alpha=.4,linestyle="dashed")
-    axs[4].plot(np.linalg.norm(r[:,2:4],axis=1),c="black",alpha=.4,linestyle="dashed")
+for r in reproductions:
+    axs[0].plot(r[:, 0], c="black", alpha=.4, linestyle="dashed")
+    axs[1].plot(r[:, 1], c="black", alpha=.4, linestyle="dashed")
+    axs[2].plot(r[:, 2], c="black", alpha=.4, linestyle="dashed")
+    axs[3].plot(r[:, 3], c="black", alpha=.4, linestyle="dashed")
+    axs[4].plot(np.linalg.norm(r[:, 2:4], axis=1), c="black", alpha=.4, linestyle="dashed")
 
 axs[0].set_ylabel("$x_1$")
 axs[1].set_ylabel("$x_2$")