locationd: make pose_kf inherit from KalmanFilter (#34982)

* Read message not json for initial state

* Delete lines

* Fix param

selfdrive/locationd/locationd.py

@@ -1,6 +1,5 @@
 #!/usr/bin/env python3
 import os
-import json
 import time
 import capnp
 import numpy as np
@@ -65,7 +64,7 @@ class LocationEstimator:
     self.observation_errors = {kind: np.zeros(3, dtype=np.float32) for kind in obs_kinds}
 
   def reset(self, t: float, x_initial: np.ndarray = PoseKalman.initial_x, P_initial: np.ndarray = PoseKalman.initial_P):
-    self.kf.reset(t, x_initial, P_initial)
+    self.kf.init_state(x_initial, covs=P_initial, filter_time=t)
 
   def _validate_sensor_source(self, source: log.SensorEventData.SensorSource):
     # some segments have two IMUs, ignore the second one
@@ -186,8 +185,8 @@ class LocationEstimator:
       rot_device_noise = rot_device_std ** 2
       trans_device_noise = trans_device_std ** 2
 
-      cam_odo_rot_res = self.kf.predict_and_observe(t, ObservationKind.CAMERA_ODO_ROTATION, rot_device, rot_device_noise)
-      cam_odo_trans_res = self.kf.predict_and_observe(t, ObservationKind.CAMERA_ODO_TRANSLATION, trans_device, trans_device_noise)
+      cam_odo_rot_res = self.kf.predict_and_observe(t, ObservationKind.CAMERA_ODO_ROTATION, rot_device, np.array([np.diag(rot_device_noise)]))
+      cam_odo_trans_res = self.kf.predict_and_observe(t, ObservationKind.CAMERA_ODO_TRANSLATION, trans_device, np.array([np.diag(trans_device_noise)]))
       self.camodo_yawrate_distribution =  np.array([rot_device[2], rot_device_std[2]])
       if cam_odo_rot_res is not None:
         _, new_x, _, new_P, _, _, (cam_odo_rot_err,), _, _ = cam_odo_rot_res
@@ -278,12 +277,13 @@ def main():
   input_invalid_threshold = {s: input_invalid_limit[s] - 0.5 for s in critcal_services}
   input_invalid_decay = {s: calculate_invalid_input_decay(input_invalid_limit[s], INPUT_INVALID_RECOVERY, SERVICE_LIST[s].frequency) for s in critcal_services}
 
-  initial_pose = params.get("LocationFilterInitialState")
-  if initial_pose is not None:
-    initial_pose = json.loads(initial_pose)
-    x_initial = np.array(initial_pose["x"], dtype=np.float64)
-    P_initial = np.diag(np.array(initial_pose["P"], dtype=np.float64))
-    estimator.reset(None, x_initial, P_initial)
+  initial_pose_data = params.get("LocationFilterInitialState")
+  if initial_pose_data is not None:
+    with log.Event.from_bytes(initial_pose_data) as lp_msg:
+      filter_state = lp_msg.livePose.debugFilterState
+      x_initial = np.array(filter_state.value, dtype=np.float64) if len(filter_state.value) != 0 else PoseKalman.initial_x
+      P_initial = np.diag(np.array(filter_state.std, dtype=np.float64)) if len(filter_state.std) != 0 else PoseKalman.initial_P
+      estimator.reset(None, x_initial, P_initial)
 
   while True:
     sm.update()

selfdrive/locationd/models/pose_kf.py

@@ -5,6 +5,8 @@ import numpy as np
 
 from openpilot.selfdrive.locationd.models.constants import ObservationKind
 
+from rednose.helpers.kalmanfilter import KalmanFilter
+
 if __name__=="__main__":
   import sympy as sp
   from rednose.helpers.ekf_sym import gen_code
@@ -24,7 +26,7 @@ class States:
   ACCEL_BIAS = slice(15, 18)  # Acceletometer bias in m/s**2
 
 
-class PoseKalman:
+class PoseKalman(KalmanFilter):
   name = "pose"
 
   # state
@@ -50,10 +52,10 @@ class PoseKalman:
                3**2, 3**2, 3**2,
                0.005**2, 0.005**2, 0.005**2])
 
-  obs_noise = {ObservationKind.PHONE_GYRO: np.array([0.025**2, 0.025**2, 0.025**2]),
-               ObservationKind.PHONE_ACCEL: np.array([.5**2, .5**2, .5**2]),
-               ObservationKind.CAMERA_ODO_TRANSLATION: np.array([0.5**2, 0.5**2, 0.5**2]),
-               ObservationKind.CAMERA_ODO_ROTATION: np.array([0.05**2, 0.05**2, 0.05**2])}
+  obs_noise = {ObservationKind.PHONE_GYRO: np.diag([0.025**2, 0.025**2, 0.025**2]),
+               ObservationKind.PHONE_ACCEL: np.diag([.5**2, .5**2, .5**2]),
+               ObservationKind.CAMERA_ODO_TRANSLATION: np.diag([0.5**2, 0.5**2, 0.5**2]),
+               ObservationKind.CAMERA_ODO_ROTATION: np.diag([0.05**2, 0.05**2, 0.05**2])}
 
   @staticmethod
   def generate_code(generated_dir):
@@ -103,35 +105,6 @@ class PoseKalman:
     self.filter = EKF_sym_pyx(generated_dir, self.name, PoseKalman.Q, PoseKalman.initial_x, PoseKalman.initial_P,
                               dim_state, dim_state_err, max_rewind_age=max_rewind_age)
 
-  @property
-  def x(self):
-    return self.filter.state()
-
-  @property
-  def P(self):
-    return self.filter.covs()
-
-  @property
-  def t(self):
-    return self.filter.get_filter_time()
-
-  def predict_and_observe(self, t, kind, data, obs_noise=None):
-    data = np.atleast_2d(data)
-    if obs_noise is None:
-      obs_noise = self.obs_noise[kind]
-    R = self._get_R(len(data), obs_noise)
-    return self.filter.predict_and_update_batch(t, kind, data, R)
-
-  def reset(self, t, x_init, P_init):
-    self.filter.init_state(x_init, P_init, t)
-
-  def _get_R(self, n, obs_noise):
-    dim = obs_noise.shape[0]
-    R = np.zeros((n, dim, dim))
-    for i in range(n):
-      R[i, :, :] = np.diag(obs_noise)
-    return R
-
 
 if __name__ == "__main__":
   generated_dir = sys.argv[2]