fix: modified LQT_recursive_LS examples to be closer from matlab examples

python/LQT_recursive_LS.py

-'''
+"""
 Linear Quadratic tracker applied on a via point example
 
 Copyright (c) 2023 Idiap Research Institute <https://www.idiap.ch/>
@@ -7,7 +7,7 @@ Sylvain Calinon <https://calinon.ch>
 
 This file is part of RCFS <https://robotics-codes-from-scratch.github.io/>
 License: GPL-3.0-only
-'''
+"""
 
 import numpy as np
 from math import factorial
@@ -15,102 +15,135 @@ import matplotlib.pyplot as plt
 
 # Parameters
 # ===============================
-param = lambda: None # Lazy way to define an empty class in python
+param = lambda: None  # Lazy way to define an empty class in python
 param.dt = 1e-2  # Time step length
-param.nbPoints = 1 # Number of target point
-param.nbDeriv = 2 # Order of the dynamical system
-param.nbVarPos = 2 # Number of position variable
-param.nbVar = param.nbVarPos * param.nbDeriv + 1 # Dimension of state vector
+param.nbPoints = 2  # Number of target point
+param.nbDeriv = 2  # Order of the dynamical system
+param.nbVarPos = 2  # Number of position variable
+param.nbVar = param.nbVarPos * param.nbDeriv + 1  # Dimension of state vector
 param.nbData = 100  # Number of datapoints
-param.rfactor = 1e-9  # Control weight term
+param.rfactor = 1e-6  # Control weight term
+param.perturbation_id = 30
 
-R = np.eye((param.nbData-1) * param.nbVarPos) * param.rfactor  # Control cost matrix
-Q = np.zeros((param.nbVar * param.nbData, param.nbVar * param.nbData)) # Task precision for augmented state
+R = np.eye((param.nbData - 1) * param.nbVarPos) * param.rfactor  # Control cost matrix
+Q = np.zeros(
+    (param.nbVar * param.nbData, param.nbVar * param.nbData)
+)  # Task precision for augmented state
 
 via_points = []
-via_point_timing = np.linspace(0, param.nbData-1, param.nbPoints+1)
+via_point_timing = np.linspace(0, param.nbData - 1, param.nbPoints + 1)
 
-for tv_ in via_point_timing[1:]:
+# Via-points initialization
+# =====================================
+# Randomly
+# for tv_ in via_point_timing[1:]:
+#     via_points.append(np.random.uniform(size=param.nbVarPos))
+via_points.append([0.35, 0.75])
+via_points.append([0.75, 0.35])
+
+
+for via_point, tv_ in zip(via_points, via_point_timing[1:]):
     tv = int(tv_)
-    tv_slice = slice(tv*param.nbVar, (tv+1)*param.nbVar)
-    X_d_tv = np.concatenate((np.random.uniform(size=param.nbVarPos), np.zeros(param.nbVarPos)))
-    via_points.append(X_d_tv[:param.nbVarPos])
-    print("via point: ", X_d_tv)
-    Q_tv = np.eye(param.nbVar-1)
-    if tv < param.nbData-1:
-      Q_tv[param.nbVarPos:, param.nbVarPos:] = 1e-6*np.eye(param.nbVarPos) # Don't track velocities on the way
+    tv_slice = slice(tv * param.nbVar, (tv + 1) * param.nbVar)
+    X_d_tv = np.concatenate((via_point, np.zeros(param.nbVarPos)))
+    Q_tv = np.eye(param.nbVar - 1)
+    if tv < param.nbData - 1:
+        Q_tv[param.nbVarPos :, param.nbVarPos :] = 1e-6 * np.eye(
+            param.nbVarPos
+        )  # Don't track velocities on the way
     Q_inv = np.zeros((param.nbVar, param.nbVar))
-    Q_inv[:param.nbVar-1, :param.nbVar-1] = np.linalg.inv(Q_tv) + np.outer(X_d_tv, X_d_tv)
-    Q_inv[:param.nbVar-1, -1] = X_d_tv
-    Q_inv[-1, :param.nbVar-1] = X_d_tv
+    Q_inv[: param.nbVar - 1, : param.nbVar - 1] = np.linalg.inv(Q_tv) + np.outer(
+        X_d_tv, X_d_tv
+    )
+    Q_inv[: param.nbVar - 1, -1] = X_d_tv
+    Q_inv[-1, : param.nbVar - 1] = X_d_tv
     Q_inv[-1, -1] = 1
     Q[tv_slice, tv_slice] = np.linalg.inv(Q_inv)
 
 
 # Dynamical System settings (discrete)
 # =====================================
-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.eye(param.nbVar)
-A[:param.nbVar-1, :param.nbVar-1] = np.kron(A1d, np.identity(param.nbVarPos))
+A[: param.nbVar - 1, : param.nbVar - 1] = np.kron(A1d, np.identity(param.nbVarPos))
 B = np.zeros((param.nbVar, param.nbVarPos))
-B[:param.nbVar-1] = np.kron(B1d, np.identity(param.nbVarPos))
+B[: param.nbVar - 1] = np.kron(B1d, np.identity(param.nbVarPos))
 
 # Build Sx and Su transfer matrices
-Su = np.zeros((param.nbVar*param.nbData, param.nbVarPos * (param.nbData-1))) 
-Sx = np.kron(np.ones((param.nbData,1)), np.eye(param.nbVar,param.nbVar))
+Su = np.zeros((param.nbVar * param.nbData, param.nbVarPos * (param.nbData - 1)))
+Sx = np.kron(np.ones((param.nbData, 1)), np.eye(param.nbVar, param.nbVar))
 
 M = B
-for i in range(1,param.nbData):
-    Sx[i*param.nbVar:param.nbData*param.nbVar,:] = np.dot(Sx[i*param.nbVar:param.nbData*param.nbVar,:], A)
-    Su[param.nbVar*i:param.nbVar*i+M.shape[0],0:M.shape[1]] = M
-    M = np.hstack((np.dot(A,M),B)) # [0,nb_state_var-1]
+for i in range(1, param.nbData):
+    Sx[i * param.nbVar : param.nbData * param.nbVar, :] = np.dot(
+        Sx[i * param.nbVar : param.nbData * param.nbVar, :], A
+    )
+    Su[param.nbVar * i : param.nbVar * i + M.shape[0], 0 : M.shape[1]] = M
+    M = np.hstack((np.dot(A, M), B))  # [0,nb_state_var-1]
 
 
 # Build recursive least squares solution for the feedback gains, using u_k = - F_k @ x0 = - K_k @ x_k
 # =====================================
 F = -np.linalg.inv(Su.T @ Q @ Su + R) @ Su.T @ Q @ Sx
-K = np.zeros((param.nbData-1, param.nbVarPos, param.nbVar))
+K = np.zeros((param.nbData - 1, param.nbVarPos, param.nbVar))
 P = np.eye(param.nbVar)
-K[0] = F[:param.nbVarPos]
-for k in range(1, param.nbData-1):
-  k_slice = slice(k*param.nbVarPos, (k+1)*param.nbVarPos)
-  P = P @ np.linalg.inv(A + B @ K[k-1])
-  K[k] = F[k_slice] @ P
+K[0] = F[: param.nbVarPos]
+for k in range(1, param.nbData - 1):
+    k_slice = slice(k * param.nbVarPos, (k + 1) * param.nbVarPos)
+    P = P @ np.linalg.inv(A + B @ K[k - 1])
+    K[k] = F[k_slice] @ P
 
 
 # Multiple reproductions from different initial positions with perturbation
 # =====================================
-num_repro = 10
+num_repro = 1
 
 x = np.zeros((param.nbVar, num_repro))
-x[:param.nbVarPos, :] = np.random.normal(0, 1e-1, (param.nbVarPos, num_repro))
-x[-1,:] = np.ones(num_repro)
+# x[:param.nbVarPos, :] = np.random.normal(0, 1e-0, (param.nbVarPos, num_repro)) # for random x0 initialisation
+x[-1, :] = np.ones(num_repro)
 
 X_batch = []
-X_batch.append(x[:-1,:].copy())
+X_batch.append(x[:-1, :].copy())
+
+X_batch_no_perturbation = []
+x_no_pert = np.copy(x)
+
+for k in range(param.nbData - 1):
+    u = K[k] @ x
+    u_no_pert = K[k] @ x_no_pert
+    if k == param.perturbation_id:
+        # simulate a random perturbation at this timestep
+        x[0 : param.nbVarPos] += np.random.normal(0, 0.2, (param.nbVarPos, num_repro))
+    x = A @ x + B @ u
+    x_no_pert = A @ x_no_pert + B @ u_no_pert
+    X_batch.append(x[:-1, :].copy())
+    X_batch_no_perturbation.append(x_no_pert[:-1, :].copy())
 
-for k in range(param.nbData-1):
-  u = K[k] @ x
-  x = A @ x + B @ (u + np.random.normal(0, 1e+1, (param.nbVarPos, num_repro)))
-  X_batch.append(x[:-1,:].copy())
-  
 X_traj = np.array(X_batch)
+X_traj_no_pert = np.array(X_batch_no_perturbation)
+perturbation = X_traj[param.perturbation_id : param.perturbation_id + 2]
 
 # Plotting
 # =========
-plt.figure()
-plt.title("2D Trajectory")
-plt.scatter(X_traj[0,0], X_traj[0,1], c='black', s=100)
-plt.plot(X_traj[:,0], X_traj[:,1], c='black')
-for p in via_points:
-  plt.scatter(p[0], p[1], c='red', s=100)  
-plt.axis("off")
-plt.gca().set_aspect('equal', adjustable='box')
-
+fig, ax = plt.subplots()
+plt.title("Linear Quadratic tracker applied on a via point examples")
+ax.plot(X_traj[:, 0], X_traj[:, 1], marker=".", alpha=0.3)
+lines = ax.plot(perturbation[:, 0], perturbation[:, 1], c="r", marker=".")
+ax.plot(
+    X_traj_no_pert[:, 0], X_traj_no_pert[:, 1], c="black", linestyle="dashed", alpha=0.2
+)
+for p in via_points[:-1]:
+    ax.scatter(p[0], p[1], c="g", s=100, label="Via points")
+target = via_points[-1]
+ax.scatter(target[0], target[1], c="b", s=100, label="Target")
+ax.axis("off")
+ax.set_aspect("equal", adjustable="box")
+lines[0].set_label("Perturbation")
+plt.legend()
 plt.show()

python/LQT_recursive_LS_multiAgents.py

-'''
+"""
 Linear Quadratic tracker applied on a via point example while coordinating two agents
 
 Copyright (c) 2023 Idiap Research Institute <https://www.idiap.ch/>
@@ -7,7 +7,7 @@ Sylvain Calinon <https://calinon.ch>
 
 This file is part of RCFS <https://robotics-codes-from-scratch.github.io/>
 License: GPL-3.0-only
-'''
+"""
 
 import numpy as np
 from math import factorial
@@ -15,98 +15,117 @@ import matplotlib.pyplot as plt
 
 # Parameters
 # ===============================
-param = lambda: None # Lazy way to define an empty class in python
+param = lambda: None  # Lazy way to define an empty class in python
 param.dt = 1e-2  # Time step length
-param.nbPoints = 1 # Number of target points
-param.nbDeriv = 2 # Order of the dynamical system
-param.nbVarPos = 2 # Number of position variable
-param.nbVar = 2 * param.nbVarPos * param.nbDeriv + 1 # Dimension of state vector: [x_1, x_2, dx_1, dx_2, 1]
+param.nbPoints = 1  # Number of target points
+param.nbDeriv = 2  # Order of the dynamical system
+param.nbVarPos = 2  # Number of position variable
+param.nbVar = (
+    2 * param.nbVarPos * param.nbDeriv + 1
+)  # Dimension of state vector: [x_1, x_2, dx_1, dx_2, 1]
 param.nbVarU = 2 * param.nbVarPos
 param.nbData = 100  # Number of datapoints
 param.rfactor = 1e-9  # Control weight term
+param.perturbation_id = 30
 
-R = np.eye((param.nbData-1) * param.nbVarU) * param.rfactor  # Control cost matrix
-Q = np.zeros((param.nbVar * param.nbData, param.nbVar * param.nbData)) # Task precision for augmented state
+R = np.eye((param.nbData - 1) * param.nbVarU) * param.rfactor  # Control cost matrix
+Q = np.zeros(
+    (param.nbVar * param.nbData, param.nbVar * param.nbData)
+)  # Task precision for augmented state
 
 via_points = []
-via_point_timing = np.linspace(0, param.nbData-1, param.nbPoints+1)
+via_point_timing = np.linspace(0, param.nbData - 1, param.nbPoints + 1)
 
 for tv_ in via_point_timing[1:]:
     tv = int(tv_)
-    tv_slice = slice(tv*param.nbVar, (tv+1)*param.nbVar)
-    
-    X_d_tv = np.concatenate((np.random.uniform(-1.0, 1.0, size=2 * param.nbVarPos), np.zeros(2 * param.nbVarPos)))
-    
-    via_points.append(X_d_tv[:2*param.nbVarPos])
-    print("via point: ", X_d_tv)
-    
-    Q_tv = np.eye(param.nbVar-1)
-    if tv < param.nbData-1:
-      Q_tv[2*param.nbVarPos:, 2*param.nbVarPos:] = 1e-6*np.eye(2*param.nbVarPos) # Don't track velocities on the way
-    
+    tv_slice = slice(tv * param.nbVar, (tv + 1) * param.nbVar)
+
+    X_d_tv = np.array([0.75, -0.5, -0.25, -1.0, 0.0, 0.0, 0.0, 0.0])
+    # Or for random x0 initialisation
+    # X_d_tv = np.concatenate((np.random.uniform(-1.0, 1.0, size=2 * param.nbVarPos), np.zeros(2 * param.nbVarPos)))
+    via_points.append(X_d_tv[: 2 * param.nbVarPos])
+
+    Q_tv = np.eye(param.nbVar - 1)
+    if tv < param.nbData - 1:
+        Q_tv[2 * param.nbVarPos :, 2 * param.nbVarPos :] = 1e-6 * np.eye(
+            2 * param.nbVarPos
+        )  # Don't track velocities on the way
+
     Q_inv = np.zeros((param.nbVar, param.nbVar))
-    Q_inv[:param.nbVar-1, :param.nbVar-1] = np.linalg.inv(Q_tv) + np.outer(X_d_tv, X_d_tv)
-    Q_inv[:param.nbVar-1, -1] = X_d_tv
-    Q_inv[-1, :param.nbVar-1] = X_d_tv
+    Q_inv[: param.nbVar - 1, : param.nbVar - 1] = np.linalg.inv(Q_tv) + np.outer(
+        X_d_tv, X_d_tv
+    )
+    Q_inv[: param.nbVar - 1, -1] = X_d_tv
+    Q_inv[-1, : param.nbVar - 1] = X_d_tv
     Q_inv[-1, -1] = 1
-    
+
     Q[tv_slice, tv_slice] = np.linalg.inv(Q_inv)
-    
+
 # Construct meeting point constraint at half of the horizon
 tv = int(0.5 * param.nbData)
-tv_slice = slice(tv*param.nbVar, (tv+1)*param.nbVar)
-Q_tv = np.eye(param.nbVar-1)
+tv_slice = slice(tv * param.nbVar, (tv + 1) * param.nbVar)
+Q_tv = np.eye(param.nbVar - 1)
 
 # Off-diagonal to tie position
-Q_tv[:param.nbVarPos, param.nbVarPos:2*param.nbVarPos] = - (1.0 - 1e-6) * np.eye(param.nbVarPos)
-Q_tv[param.nbVarPos:2*param.nbVarPos, :param.nbVarPos] = - (1.0 - 1e-6) * np.eye(param.nbVarPos)
+Q_tv[: param.nbVarPos, param.nbVarPos : 2 * param.nbVarPos] = -(1.0 - 1e-6) * np.eye(
+    param.nbVarPos
+)
+Q_tv[param.nbVarPos : 2 * param.nbVarPos, : param.nbVarPos] = -(1.0 - 1e-6) * np.eye(
+    param.nbVarPos
+)
 
 # Off-diagonal to tie velocity
-Q_tv[2*param.nbVarPos:3*param.nbVarPos, 3*param.nbVarPos:4*param.nbVarPos] = - (1.0 - 1e-6) * np.eye(param.nbVarPos)
-Q_tv[3*param.nbVarPos:4*param.nbVarPos, 2*param.nbVarPos:3*param.nbVarPos] = - (1.0 - 1e-6) * np.eye(param.nbVarPos)
-
-#Q_tv[2*param.nbVarPos:, 2*param.nbVarPos:] = 1e-6*np.eye(2*param.nbVarPos) # Don't track absolute velocities
+Q_tv[
+    2 * param.nbVarPos : 3 * param.nbVarPos, 3 * param.nbVarPos : 4 * param.nbVarPos
+] = -(1.0 - 1e-6) * np.eye(param.nbVarPos)
+Q_tv[
+    3 * param.nbVarPos : 4 * param.nbVarPos, 2 * param.nbVarPos : 3 * param.nbVarPos
+] = -(1.0 - 1e-6) * np.eye(param.nbVarPos)
+
+# Q_tv[2*param.nbVarPos:, 2*param.nbVarPos:] = 1e-6*np.eye(2*param.nbVarPos) # Don't track absolute velocities
 Q_inv = np.zeros((param.nbVar, param.nbVar))
-Q_inv[:param.nbVar-1, :param.nbVar-1] = np.linalg.inv(Q_tv)
+Q_inv[: param.nbVar - 1, : param.nbVar - 1] = np.linalg.inv(Q_tv)
 Q_inv[-1, -1] = 1
 Q[tv_slice, tv_slice] = np.linalg.inv(Q_inv)
 
 
 # Dynamical System settings (discrete)
 # =====================================
-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.eye(param.nbVar)
-A[:param.nbVar-1, :param.nbVar-1] = np.kron(A1d, np.identity(2 * param.nbVarPos))
+A[: param.nbVar - 1, : param.nbVar - 1] = np.kron(A1d, np.identity(2 * param.nbVarPos))
 B = np.zeros((param.nbVar, 2 * param.nbVarPos))
-B[:param.nbVar-1] = np.kron(B1d, np.identity(2 * param.nbVarPos))
+B[: param.nbVar - 1] = np.kron(B1d, np.identity(2 * param.nbVarPos))
 
 # Build Sx and Su transfer matrices
-Su = np.zeros((param.nbVar*param.nbData, param.nbVarU * (param.nbData-1))) 
-Sx = np.kron(np.ones((param.nbData,1)), np.eye(param.nbVar,param.nbVar))
+Su = np.zeros((param.nbVar * param.nbData, param.nbVarU * (param.nbData - 1)))
+Sx = np.kron(np.ones((param.nbData, 1)), np.eye(param.nbVar, param.nbVar))
 
 M = B
-for i in range(1,param.nbData):
-    Sx[i*param.nbVar:param.nbData*param.nbVar,:] = np.dot(Sx[i*param.nbVar:param.nbData*param.nbVar,:], A)
-    Su[param.nbVar*i:param.nbVar*i+M.shape[0],0:M.shape[1]] = M
-    M = np.hstack((np.dot(A,M),B)) # [0,nb_state_var-1]
+for i in range(1, param.nbData):
+    Sx[i * param.nbVar : param.nbData * param.nbVar, :] = np.dot(
+        Sx[i * param.nbVar : param.nbData * param.nbVar, :], A
+    )
+    Su[param.nbVar * i : param.nbVar * i + M.shape[0], 0 : M.shape[1]] = M
+    M = np.hstack((np.dot(A, M), B))  # [0,nb_state_var-1]
 
 
 # Build recursive least squares solution for the feedback gains, using u_k = - F_k @ x0 = - K_k @ x_k
 # =====================================
 F = -np.linalg.inv(Su.T @ Q @ Su + R) @ Su.T @ Q @ Sx
-K = np.zeros((param.nbData-1, param.nbVarU, param.nbVar))
+K = np.zeros((param.nbData - 1, param.nbVarU, param.nbVar))
 P = np.eye(param.nbVar)
-K[0] = F[:param.nbVarU]
-for k in range(1, param.nbData-1):
-  k_slice = slice(k*param.nbVarU, (k+1)*param.nbVarU)
-  P = P @ np.linalg.inv(A + B @ K[k-1])
-  K[k] = F[k_slice] @ P
+K[0] = F[: param.nbVarU]
+for k in range(1, param.nbData - 1):
+    k_slice = slice(k * param.nbVarU, (k + 1) * param.nbVarU)
+    P = P @ np.linalg.inv(A + B @ K[k - 1])
+    K[k] = F[k_slice] @ P
 
 
 # Multiple reproductions from different initial positions with perturbation
@@ -114,38 +133,73 @@ for k in range(1, param.nbData-1):
 num_repro = 1
 
 x = np.zeros((param.nbVar, num_repro))
-x[:param.nbVarPos, :] = np.random.normal(-0.5, 1e-1, (param.nbVarPos, num_repro))
-x[param.nbVarPos:2*param.nbVarPos, :] = np.random.normal(0.5, 1e-1, (param.nbVarPos, num_repro))
-x[-1,:] = np.ones(num_repro)
+x[: param.nbVarPos, :] = np.array([-1.0, 0.0]).reshape(-1, 1)
+x[param.nbVarPos : 2 * param.nbVarPos, :] = np.array([0.0, 1.0]).reshape(-1, 1)
+# Or for random x0 initialisation
+# x[:param.nbVarPos, :] = np.random.normal(-0.5, 1e-1, (param.nbVarPos, num_repro))
+# x[param.nbVarPos:2*param.nbVarPos, :] = np.random.normal(0.5, 1e-1, (param.nbVarPos, num_repro))
+x[-1, :] = np.ones(num_repro)
 
 X_batch = []
-X_batch.append(x[:-1,:].copy())
+X_batch.append(x[:-1, :].copy())
 
-for k in range(param.nbData-1):
-  u = K[k] @ x
-  x = A @ x + B @ (u + np.random.normal(0, 5e+0, (param.nbVarU, num_repro)))
-  X_batch.append(x[:-1,:].copy())
-  
-X_traj = np.array(X_batch)
+X_batch_no_perturbation = []
+x_no_pert = np.copy(x)
+
+for k in range(param.nbData - 1):
+    u_no_pert = K[k] @ x_no_pert
+    u = K[k] @ x
+    if k == param.perturbation_id:
+        x[0 : param.nbVarPos] += np.random.normal(0, 0.2, (param.nbVarPos, num_repro))
+    x = A @ x + B @ u
+    x_no_pert = A @ x_no_pert + B @ u_no_pert
+    X_batch.append(x[:-1, :].copy())
+    X_batch_no_perturbation.append(x_no_pert[:-1, :].copy())
 
+X_traj = np.array(X_batch)
+X_traj_no_pert = np.array(X_batch_no_perturbation)
+perturbation = X_traj[param.perturbation_id : param.perturbation_id + 2]
 
 # Plotting
 # =========
-plt.figure()
-plt.title("2D Trajectory")
+fig, ax = plt.subplots()
+plt.title(
+    "Linear Quadratic tracker applied on a via point example \n while coordinating two agents"
+)
+
 # Agent 1
-plt.scatter(X_traj[0,0], X_traj[0,1], c='black', s=100)
-plt.plot(X_traj[:,0], X_traj[:,1], c='black')
+ax.plot(X_traj[:, 0], X_traj[:, 1], marker=".", alpha=0.3)
+lines = ax.plot(perturbation[:, 0], perturbation[:, 1], c="r", marker=".")
+ax.plot(
+    X_traj_no_pert[:, 0], X_traj_no_pert[:, 1], c="black", linestyle="dashed", alpha=0.2
+)
 for p in via_points:
-  plt.scatter(p[0], p[1], c='red', s=100)
-  
+    ax.scatter(p[0], p[1], c="b", s=100)
+ax.scatter(
+    X_traj[tv][0],
+    X_traj[tv][1],
+    c="k",
+    marker="x",
+    s=100,
+    alpha=1.0,
+    label="Meeting point at half horizon",
+)
+ax.scatter(
+    X_traj_no_pert[tv][0], X_traj_no_pert[tv][1], c="k", marker="x", s=100, alpha=0.5
+)
+target = via_points[-1]
+ax.scatter(target[0], target[1], c="b", s=100, label="Target 1")
+
 # Agent 2
-plt.scatter(X_traj[0,2], X_traj[0,3], c='black', s=100)
-plt.plot(X_traj[:,2], X_traj[:,3], c='black')
-for p in via_points:
-  plt.scatter(p[2], p[3], c='red', s=100)  
+ax.plot(X_traj[:, 2], X_traj[:, 3], marker=".", alpha=0.3)
+ax.plot(
+    X_traj_no_pert[:, 2], X_traj_no_pert[:, 3], c="black", linestyle="dashed", alpha=0.2
+)
+ax.scatter(target[2], target[3], c="r", s=100, label="Target 2")
 
 plt.axis("off")
-plt.gca().set_aspect('equal', adjustable='box')
+plt.gca().set_aspect("equal", adjustable="box")
+lines[0].set_label("Perturbation")
+plt.legend()
 
 plt.show()