From 300a07de1979c5788dec88889e34496bfabbad2b Mon Sep 17 00:00:00 2001
From: FrogAi <91348155+FrogAi@users.noreply.github.com>
Date: Wed, 26 Jun 2024 01:49:39 -0700
Subject: [PATCH] Controls - Driving Personalities - Customize Personalities

Customize the driving personality profiles to your driving style.
---
 .../lib/longitudinal_mpc_lib/long_mpc.py      | 55 +++++++++++++------
 .../frogpilot/controls/frogpilot_planner.py   | 15 ++++-
 2 files changed, 52 insertions(+), 18 deletions(-)
 mode change 100755 => 100644 selfdrive/controls/lib/longitudinal_mpc_lib/long_mpc.py

diff --git a/selfdrive/controls/lib/longitudinal_mpc_lib/long_mpc.py b/selfdrive/controls/lib/longitudinal_mpc_lib/long_mpc.py
old mode 100755
new mode 100644
index cadbf68bc..bb80efde2
--- a/selfdrive/controls/lib/longitudinal_mpc_lib/long_mpc.py
+++ b/selfdrive/controls/lib/longitudinal_mpc_lib/long_mpc.py
@@ -57,26 +57,49 @@ T_DIFFS = np.diff(T_IDXS, prepend=[0.])
 COMFORT_BRAKE = 2.5
 STOP_DISTANCE = 6.0
 
-def get_jerk_factor(personality=log.LongitudinalPersonality.standard):
-  if personality==log.LongitudinalPersonality.relaxed:
-    return 1.0, 1.0, 1.0
-  elif personality==log.LongitudinalPersonality.standard:
-    return 1.0, 1.0, 1.0
-  elif personality==log.LongitudinalPersonality.aggressive:
-    return 0.5, 1.0, 0.5
+def get_jerk_factor(aggressive_jerk_acceleration=0.5, aggressive_jerk_danger=0.5, aggressive_jerk_speed=0.5,
+                    standard_jerk_acceleration=1.0, standard_jerk_danger=1.0, standard_jerk_speed=1.0,
+                    relaxed_jerk_acceleration=1.0, relaxed_jerk_danger=1.0, relaxed_jerk_speed=1.0,
+                    custom_personalities=False, personality=log.LongitudinalPersonality.standard):
+  if custom_personalities:
+    if personality==log.LongitudinalPersonality.relaxed:
+      return relaxed_jerk_acceleration, relaxed_jerk_danger, relaxed_jerk_speed
+    elif personality==log.LongitudinalPersonality.standard:
+      return standard_jerk_acceleration, standard_jerk_danger, standard_jerk_speed
+    elif personality==log.LongitudinalPersonality.aggressive:
+      return aggressive_jerk_acceleration, aggressive_jerk_danger, aggressive_jerk_speed
+    else:
+      raise NotImplementedError("Longitudinal personality not supported")
   else:
-    raise NotImplementedError("Longitudinal personality not supported")
+    if personality==log.LongitudinalPersonality.relaxed:
+      return 1.0, 1.0, 1.0
+    elif personality==log.LongitudinalPersonality.standard:
+      return 1.0, 1.0, 1.0
+    elif personality==log.LongitudinalPersonality.aggressive:
+      return 0.5, 0.5, 0.5
+    else:
+      raise NotImplementedError("Longitudinal personality not supported")
 
 
-def get_T_FOLLOW(personality=log.LongitudinalPersonality.standard):
-  if personality==log.LongitudinalPersonality.relaxed:
-    return 1.75
-  elif personality==log.LongitudinalPersonality.standard:
-    return 1.45
-  elif personality==log.LongitudinalPersonality.aggressive:
-    return 1.25
+def get_T_FOLLOW(aggressive_follow=1.25, standard_follow=1.45, relaxed_follow=1.75, custom_personalities=False, personality=log.LongitudinalPersonality.standard):
+  if custom_personalities:
+    if personality==log.LongitudinalPersonality.relaxed:
+      return relaxed_follow
+    elif personality==log.LongitudinalPersonality.standard:
+      return standard_follow
+    elif personality==log.LongitudinalPersonality.aggressive:
+      return aggressive_follow
+    else:
+      raise NotImplementedError("Longitudinal personality not supported")
   else:
-    raise NotImplementedError("Longitudinal personality not supported")
+    if personality==log.LongitudinalPersonality.relaxed:
+      return 1.75
+    elif personality==log.LongitudinalPersonality.standard:
+      return 1.45
+    elif personality==log.LongitudinalPersonality.aggressive:
+      return 1.25
+    else:
+      raise NotImplementedError("Longitudinal personality not supported")
 
 def get_stopped_equivalence_factor(v_lead):
   return (v_lead**2) / (2 * COMFORT_BRAKE)
diff --git a/selfdrive/frogpilot/controls/frogpilot_planner.py b/selfdrive/frogpilot/controls/frogpilot_planner.py
index 71f0d2bff..61653e593 100644
--- a/selfdrive/frogpilot/controls/frogpilot_planner.py
+++ b/selfdrive/frogpilot/controls/frogpilot_planner.py
@@ -82,8 +82,19 @@ class FrogPilotPlanner:
       self.min_accel = A_CRUISE_MIN
 
   def set_follow_values(self, controlsState, frogpilotCarState, lead_distance, stopping_distance, v_ego, v_lead, frogpilot_toggles):
-    self.base_acceleration_jerk, self.base_danger_jerk, self.base_speed_jerk = get_jerk_factor(controlsState.personality)
-    self.t_follow = get_T_FOLLOW(controlsState.personality)
+    self.base_acceleration_jerk, self.base_danger_jerk, self.base_speed_jerk = get_jerk_factor(
+      frogpilot_toggles.aggressive_jerk_acceleration, frogpilot_toggles.aggressive_jerk_danger, frogpilot_toggles.aggressive_jerk_speed,
+      frogpilot_toggles.standard_jerk_acceleration, frogpilot_toggles.standard_jerk_danger, frogpilot_toggles.standard_jerk_speed,
+      frogpilot_toggles.relaxed_jerk_acceleration, frogpilot_toggles.relaxed_jerk_danger, frogpilot_toggles.relaxed_jerk_speed,
+      frogpilot_toggles.custom_personalities, controlsState.personality
+    )
+
+    self.t_follow = get_T_FOLLOW(
+      frogpilot_toggles.aggressive_follow,
+      frogpilot_toggles.standard_follow,
+      frogpilot_toggles.relaxed_follow,
+      frogpilot_toggles.custom_personalities, controlsState.personality
+    )
 
     if self.tracking_lead:
       self.update_follow_values(lead_distance, stopping_distance, v_ego, v_lead, frogpilot_toggles)