From c1c40321b66445800bc2dd74306dba8d90b3d16c Mon Sep 17 00:00:00 2001
From: dragonpilot <dgnpilot@gmail.com>
Date: Fri, 27 Sep 2019 15:11:29 +1000
Subject: [PATCH] Add zorrobyte's curvature learner

---
 selfdrive/controls/lib/curvature_learner.py | 66 +++++++++++++++++++++
 selfdrive/controls/lib/pathplanner.py       | 11 +++-
 2 files changed, 76 insertions(+), 1 deletion(-)
 create mode 100644 selfdrive/controls/lib/curvature_learner.py

diff --git a/selfdrive/controls/lib/curvature_learner.py b/selfdrive/controls/lib/curvature_learner.py
new file mode 100644
index 000000000..507aed48c
--- /dev/null
+++ b/selfdrive/controls/lib/curvature_learner.py
@@ -0,0 +1,66 @@
+from numpy import clip
+import pickle
+import csv
+import os
+
+# HOW TO
+# import this module to where you want to use it, such as from ```selfdrive.controls.lib.curvature_learner import CurvatureLearner```
+# create the object ```self.curvature_offset = CurvatureLearner(debug=False)```
+# call the update method ```self.curvature_offset.update(angle_steers - angle_offset, self.LP.d_poly)```
+# The learned curvature offsets will save and load automatically
+# If you still need help, check out how I have it implemented in the devel_curvaturefactorlearner branch
+# by Zorrobyte
+# version 2
+
+class CurvatureLearner:
+    def __init__(self, debug=False):
+        self.offset = 0.
+        self.learning_rate = 12000
+        self.frame = 0
+        self.debug = debug
+        try:
+            self.learned_offsets = pickle.load(open("/data/curvaturev2.p", "rb"))
+        except (OSError, IOError) as e:
+            self.learned_offsets = {
+                "center": 0.,
+                "leftinner": 0.,
+                "rightinner": 0.,
+                "leftouter": 0.,
+                "rightouter": 0.
+            }
+            pickle.dump(self.learned_offsets, open("/data/curvaturev2.p", "wb"))
+            os.chmod("/data/curvaturev2.p", 0o777)
+
+    def update(self, angle_steers=0., d_poly=None, v_ego=0.):
+        if angle_steers > 0.1:
+            if abs(angle_steers) < 2.:
+                self.learned_offsets["center"] -= d_poly[3] / self.learning_rate
+                self.offset = self.learned_offsets["center"]
+            elif 2. < abs(angle_steers) < 5.:
+                self.learned_offsets["leftinner"] -= d_poly[3] / self.learning_rate
+                self.offset = self.learned_offsets["leftinner"]
+            elif abs(angle_steers) > 5.:
+                self.learned_offsets["leftouter"] -= d_poly[3] / self.learning_rate
+                self.offset = self.learned_offsets["leftouter"]
+        elif angle_steers < -0.1:
+            if abs(angle_steers) < 2.:
+                self.learned_offsets["center"] += d_poly[3] / self.learning_rate
+                self.offset = self.learned_offsets["center"]
+            elif 2. < abs(angle_steers) < 5.:
+                self.learned_offsets["rightinner"] += d_poly[3] / self.learning_rate
+                self.offset = self.learned_offsets["rightinner"]
+            elif abs(angle_steers) > 5.:
+                self.learned_offsets["rightouter"] += d_poly[3] / self.learning_rate
+                self.offset = self.learned_offsets["rightouter"]
+
+        self.offset = clip(self.offset, -0.3, 0.3)
+        self.frame += 1
+
+        if self.frame == 12000:  # every 2 mins
+            pickle.dump(self.learned_offsets, open("/data/curvaturev2.p", "wb"))
+            self.frame = 0
+        if self.debug:
+            with open('/data/curvdebug.csv', 'a') as csv_file:
+                csv_file_writer = csv.writer(csv_file, delimiter=',', quotechar='"', quoting=csv.QUOTE_MINIMAL)
+                csv_file_writer.writerow([self.learned_offsets, v_ego])
+        return self.offset
diff --git a/selfdrive/controls/lib/pathplanner.py b/selfdrive/controls/lib/pathplanner.py
index b62d9f110..e1dad095c 100644
--- a/selfdrive/controls/lib/pathplanner.py
+++ b/selfdrive/controls/lib/pathplanner.py
@@ -9,6 +9,8 @@ from selfdrive.controls.lib.lateral_mpc import libmpc_py
 from selfdrive.controls.lib.drive_helpers import MPC_COST_LAT
 from selfdrive.controls.lib.lane_planner import LanePlanner
 import selfdrive.messaging as messaging
+from selfdrive.controls.lib.curvature_learner import CurvatureLearner
+
 
 LOG_MPC = os.environ.get('LOG_MPC', False)
 
@@ -28,6 +30,7 @@ class PathPlanner(object):
     self.setup_mpc(CP.steerRateCost)
     self.solution_invalid_cnt = 0
     self.path_offset_i = 0.0
+    self.curvature_offset = CurvatureLearner(debug=False)
 
   def setup_mpc(self, steer_rate_cost):
     self.libmpc = libmpc_py.libmpc
@@ -57,7 +60,13 @@ class PathPlanner(object):
     # Run MPC
     self.angle_steers_des_prev = self.angle_steers_des_mpc
     VM.update_params(sm['liveParameters'].stiffnessFactor, sm['liveParameters'].steerRatio)
-    curvature_factor = VM.curvature_factor(v_ego)
+    # curvature_factor = VM.curvature_factor(v_ego)
+    if active:
+      curvfac = self.curvature_offset.update(angle_steers - angle_offset, self.LP.d_poly, v_ego)
+    else:
+      curvfac = 0.
+
+    curvature_factor = VM.curvature_factor(v_ego) + curvfac
 
     # TODO: Check for active, override, and saturation
     # if active: