Unit_F/ekf_prediction_and_correction.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from lego_robot import *\n",
    "from math import sin, cos, pi, atan2, sqrt\n",
    "from numpy import *\n",
    "from slam_f_library import get_observations, write_cylinders,\\\n",
    "    write_error_ellipses"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class ExtendedKalmanFilterSLAM:\n",
    "    def __init__(self, state, covariance,\n",
    "                 robot_width, scanner_displacement,\n",
    "                 control_motion_factor, control_turn_factor,\n",
    "                 measurement_distance_stddev, measurement_angle_stddev):\n",
    "        # The state. This is the core data of the Kalman filter.\n",
    "        self.state = state\n",
    "        self.covariance = covariance\n",
    "\n",
    "        # Some constants.\n",
    "        self.robot_width = robot_width\n",
    "        self.scanner_displacement = scanner_displacement\n",
    "        self.control_motion_factor = control_motion_factor\n",
    "        self.control_turn_factor = control_turn_factor\n",
    "        self.measurement_distance_stddev = measurement_distance_stddev\n",
    "        self.measurement_angle_stddev = measurement_angle_stddev\n",
    "\n",
    "        # Currently, the number of landmarks is zero.\n",
    "        self.number_of_landmarks = 0\n",
    "\n",
    "    @staticmethod\n",
    "    def g(state, control, w):\n",
    "        x, y, theta = state\n",
    "        l, r = control\n",
    "        if r != l:\n",
    "            alpha = (r - l) / w\n",
    "            rad = l / alpha\n",
    "            g1 = x + (rad + w / 2.) * (sin(theta + alpha) - sin(theta))\n",
    "            g2 = y + (rad + w / 2.) * (-cos(theta + alpha) + cos(theta))\n",
    "            g3 = (theta + alpha + pi) % (2 * pi) - pi\n",
    "        else:\n",
    "            g1 = x + l * cos(theta)\n",
    "            g2 = y + l * sin(theta)\n",
    "            g3 = theta\n",
    "\n",
    "        return array([g1, g2, g3])\n",
    "\n",
    "    @staticmethod\n",
    "    def dg_dstate(state, control, w):\n",
    "        theta = state[2]\n",
    "        l, r = control\n",
    "        if r != l:\n",
    "            alpha = (r - l) / w\n",
    "            theta_ = theta + alpha\n",
    "            rpw2 = l / alpha + w / 2.0\n",
    "            m = array([[1.0, 0.0, rpw2 * (cos(theta_) - cos(theta))],\n",
    "                       [0.0, 1.0, rpw2 * (sin(theta_) - sin(theta))],\n",
    "                       [0.0, 0.0, 1.0]])\n",
    "        else:\n",
    "            m = array([[1.0, 0.0, -l * sin(theta)],\n",
    "                       [0.0, 1.0,  l * cos(theta)],\n",
    "                       [0.0, 0.0,  1.0]])\n",
    "        return m\n",
    "\n",
    "    @staticmethod\n",
    "    def dg_dcontrol(state, control, w):\n",
    "        theta = state[2]\n",
    "        l, r = tuple(control)\n",
    "        if r != l:\n",
    "            rml = r - l\n",
    "            rml2 = rml * rml\n",
    "            theta_ = theta + rml / w\n",
    "            dg1dl = w * r / rml2 * (sin(theta_) - sin(theta)) - (r + l) / (2 * rml) * cos(theta_)\n",
    "            dg2dl = w * r / rml2 * (-cos(theta_) + cos(theta)) - (r + l) / (2 * rml) * sin(theta_)\n",
    "            dg1dr = (-w * l) / rml2 * (sin(theta_) - sin(theta)) + (r + l) / (2 * rml) * cos(theta_)\n",
    "            dg2dr = (-w * l) / rml2 * (-cos(theta_) + cos(theta)) + (r + l) / (2 * rml) * sin(theta_)\n",
    "\n",
    "        else:\n",
    "            dg1dl = 0.5 * (cos(theta) + l / w * sin(theta))\n",
    "            dg2dl = 0.5 * (sin(theta) - l / w * cos(theta))\n",
    "            dg1dr = 0.5 * (-l / w * sin(theta) + cos(theta))\n",
    "            dg2dr = 0.5 * (l / w * cos(theta) + sin(theta))\n",
    "\n",
    "        dg3dl = -1.0 / w\n",
    "        dg3dr = 1.0 / w\n",
    "        m = array([[dg1dl, dg1dr], [dg2dl, dg2dr], [dg3dl, dg3dr]])\n",
    "\n",
    "        return m\n",
    "\n",
    "    def predict(self, control):\n",
    "        \"\"\"The prediction step of the Kalman filter.\"\"\"\n",
    "\n",
    "        G3 = self.dg_dstate(self.state, control, self.robot_width)\n",
    "\n",
    "        left, right = control\n",
    "        left_var = (self.control_motion_factor * left)**2 + (self.control_turn_factor * (left - right))**2\n",
    "        right_var = (self.control_motion_factor * right)**2 + (self.control_turn_factor * (left - right))**2\n",
    "        control_covariance = diag([left_var, right_var])\n",
    "        V = self.dg_dcontrol(self.state, control, self.robot_width)\n",
    "        R3 = dot(V, dot(control_covariance, V.T))\n",
    "\n",
    "\n",
    "        new_offset = 2 * self.number_of_landmarks\n",
    "        G = zeros((3 + new_offset, 3 + new_offset))\n",
    "        G[0:3, 0:3] = G3\n",
    "        G[3:, 3:] = eye(new_offset, new_offset)\n",
    "\n",
    "        R = zeros((3 + new_offset, 3 + new_offset))\n",
    "        R[0:3, 0:3] = R3\n",
    "\n",
    "        self.covariance = dot(G, dot(self.covariance, G.T)) + R\n",
    "\n",
    "\n",
    "        self.state[0:3] = self.g(self.state[0:3], control, self.robot_width)\n",
    "\n",
    "    def add_landmark_to_state(self, initial_coords):\n",
    "        \"\"\"Enlarge the current state and covariance matrix to include one more\n",
    "           landmark, which is given by its initial_coords (an (x, y) tuple).\n",
    "           Returns the index of the newly added landmark.\"\"\"\n",
    "\n",
    " \n",
    "        landmark_index = self.number_of_landmarks\n",
    "        old_offset = 2 * self.number_of_landmarks\n",
    "        self.number_of_landmarks += 1\n",
    "        new_offset = 2 * self.number_of_landmarks\n",
    "\n",
    "        covarianceN = zeros((3 + new_offset, 3 + new_offset))\n",
    "        covarianceN[0:3 + old_offset, 0:3 + old_offset] = self.covariance\n",
    "        covarianceN[-2, -2] = covarianceN[-1, -1] = 1e10\n",
    "        self.covariance = covarianceN\n",
    "\n",
    "        state_prime = zeros(3 + new_offset)\n",
    "        state_prime[0:3 + old_offset] = self.state\n",
    "        state_prime[-2] = initial_coords[0]\n",
    "        state_prime[-1] = initial_coords[1]\n",
    "        self.state = state_prime\n",
    "\n",
    "        return landmark_index\n",
    "\n",
    "    @staticmethod\n",
    "    def h(state, landmark, scanner_displacement):\n",
    "        \"\"\"Takes a (x, y, theta) state and a (x, y) landmark, and returns the\n",
    "           measurement (range, bearing).\"\"\"\n",
    "        dx = landmark[0] - (state[0] + scanner_displacement * cos(state[2]))\n",
    "        dy = landmark[1] - (state[1] + scanner_displacement * sin(state[2]))\n",
    "        r = sqrt(dx * dx + dy * dy)\n",
    "        # alpha in [-pi, pi]\n",
    "        alpha = (atan2(dy, dx) - state[2] + pi) % (2 * pi) - pi\n",
    "\n",
    "        return array([r, alpha])\n",
    "\n",
    "    @staticmethod\n",
    "    def dh_dstate(state, landmark, scanner_displacement):\n",
    "        theta = state[2]\n",
    "        cost, sint = cos(theta), sin(theta)\n",
    "        dx = landmark[0] - (state[0] + scanner_displacement * cost)\n",
    "        dy = landmark[1] - (state[1] + scanner_displacement * sint)\n",
    "        q = dx * dx + dy * dy\n",
    "        sqrtq = sqrt(q)\n",
    "        drdx = -dx / sqrtq\n",
    "        drdy = -dy / sqrtq\n",
    "        drdtheta = (dx * sint - dy * cost) * scanner_displacement / sqrtq\n",
    "        dalphadx = dy / q\n",
    "        dalphady = -dx / q\n",
    "        dalphadtheta = -1 - scanner_displacement / q * (dx * cost + dy * sint)\n",
    "\n",
    "        return array([[drdx, drdy, drdtheta],\n",
    "                      [dalphadx, dalphady, dalphadtheta]])\n",
    "\n",
    "    def correct(self, measurement, landmark_index):\n",
    "        \"\"\"The correction step of the Kalman filter.\"\"\"\n",
    "        \n",
    "        index = 3 + 2 * landmark_index\n",
    "        landmark = self.state[index: index + 2]\n",
    "        H3 = self.dh_dstate(self.state, landmark, self.scanner_displacement)\n",
    "\n",
    "        \n",
    "        H = zeros((2, 3 + 2 * self.number_of_landmarks))\n",
    "        H[:, :3] = H3\n",
    "        H[:, index:index + 2] = -H3[:, 0:2]\n",
    "\n",
    "        \n",
    "        Q = diag([self.measurement_distance_stddev**2, self.measurement_angle_stddev**2])\n",
    "        K = dot(self.covariance, dot(H.T, linalg.inv(dot(H, dot(self.covariance, H.T)) + Q)))\n",
    "        innovation = array(measurement) - self.h(self.state, landmark, self.scanner_displacement)\n",
    "       \n",
    "        innovation[1] = (innovation[1] + pi) % (2 * pi) - pi\n",
    "        self.state = self.state + dot(K, innovation)\n",
    "        self.covariance = dot(eye(size(self.state)) - dot(K, H), self.covariance)\n",
    "\n",
    "    def get_landmarks(self):\n",
    "        \"\"\"Returns a list of (x, y) tuples of all landmark positions.\"\"\"\n",
    "        return ([(self.state[3 + 2 * j], self.state[3 + 2 * j + 1])\n",
    "                 for j in xrange(self.number_of_landmarks)])\n",
    "\n",
    "    def get_landmark_error_ellipses(self):\n",
    "        \"\"\"Returns a list of all error ellipses, one for each landmark.\"\"\"\n",
    "        ellipses = []\n",
    "        for i in xrange(self.number_of_landmarks):\n",
    "            j = 3 + 2 * i\n",
    "            ellipses.append(self.get_error_ellipse(\n",
    "                self.covariance[j:j + 2, j:j + 2]))\n",
    "        return ellipses\n",
    "\n",
    "    @staticmethod\n",
    "    def get_error_ellipse(covariance):\n",
    "        \"\"\"Return the position covariance (which is the upper 2x2 submatrix)\n",
    "           as a triple: (main_axis_angle, stddev_1, stddev_2), where\n",
    "           main_axis_angle is the angle (pointing direction) of the main axis,\n",
    "           along which the standard deviation is stddev_1, and stddev_2 is the\n",
    "           standard deviation along the other (orthogonal) axis.\"\"\"\n",
    "        eigenvals, eigenvects = linalg.eig(covariance[0:2, 0:2])\n",
    "        angle = atan2(eigenvects[1, 0], eigenvects[0, 0])\n",
    "        return (angle, sqrt(eigenvals[0]), sqrt(eigenvals[1]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "if __name__ == '__main__':\n",
    "    # Robot constants.\n",
    "    scanner_displacement = 30.0\n",
    "    ticks_to_mm = 0.349\n",
    "    robot_width = 155.0\n",
    "\n",
    "    # Cylinder extraction and matching constants.\n",
    "    minimum_valid_distance = 20.0\n",
    "    depth_jump = 100.0\n",
    "    cylinder_offset = 90.0\n",
    "    max_cylinder_distance = 500.0\n",
    "\n",
    "    # Filter constants.\n",
    "    control_motion_factor = 0.35  # Error in motor control.\n",
    "    control_turn_factor = 0.6  # Additional error due to slip when turning.\n",
    "    measurement_distance_stddev = 600.0  # Distance measurement error of cylinders.\n",
    "    measurement_angle_stddev = 45. / 180.0 * pi  # Angle measurement error.\n",
    "\n",
    "    # Arbitrary start position.\n",
    "    initial_state = array([500.0, 0.0, 45.0 / 180.0 * pi])\n",
    "\n",
    "    # Covariance at start position.\n",
    "    initial_covariance = zeros((3,3))\n",
    "\n",
    "    # Setup filter.\n",
    "    kf = ExtendedKalmanFilterSLAM(initial_state, initial_covariance,\n",
    "                                  robot_width, scanner_displacement,\n",
    "                                  control_motion_factor, control_turn_factor,\n",
    "                                  measurement_distance_stddev,\n",
    "                                  measurement_angle_stddev)\n",
    "\n",
    "    # Read data.\n",
    "    logfile = LegoLogfile()\n",
    "    logfile.read(\"robot4_motors.txt\")\n",
    "    logfile.read(\"robot4_scan.txt\")\n",
    "\n",
    "    # Loop over all motor tick records and all measurements and generate\n",
    "    # filtered positions and covariances.\n",
    "    # This is the EKF SLAM loop.\n",
    "    f = open(\"ekf_slam_correction.txt\", \"w\")\n",
    "    for i in xrange(len(logfile.motor_ticks)):\n",
    "        # Prediction.\n",
    "        control = array(logfile.motor_ticks[i]) * ticks_to_mm\n",
    "        kf.predict(control)\n",
    "\n",
    "        # Correction.\n",
    "        observations = get_observations(\n",
    "            logfile.scan_data[i],\n",
    "            depth_jump, minimum_valid_distance, cylinder_offset,\n",
    "            kf, max_cylinder_distance)\n",
    "        for obs in observations:\n",
    "            measurement, cylinder_world, cylinder_scanner, cylinder_index = obs\n",
    "            if cylinder_index == -1:\n",
    "                cylinder_index = kf.add_landmark_to_state(cylinder_world)\n",
    "            kf.correct(measurement, cylinder_index)\n",
    "\n",
    "        # End of EKF SLAM - from here on, data is written.\n",
    "\n",
    "        # Output the center of the scanner, not the center of the robot.\n",
    "        print >> f, \"F %f %f %f\" % \\\n",
    "            tuple(kf.state[0:3] + [scanner_displacement * cos(kf.state[2]),\n",
    "                                   scanner_displacement * sin(kf.state[2]),\n",
    "                                   0.0])\n",
    "        # Write covariance matrix in angle stddev1 stddev2 stddev-heading form.\n",
    "        e = ExtendedKalmanFilterSLAM.get_error_ellipse(kf.covariance)\n",
    "        print >> f, \"E %f %f %f %f\" % (e + (sqrt(kf.covariance[2,2]),))\n",
    "        # Write estimates of landmarks.\n",
    "        write_cylinders(f, \"W C\", kf.get_landmarks())\n",
    "        # Write error ellipses of landmarks.\n",
    "        write_error_ellipses(f, \"W E\", kf.get_landmark_error_ellipses())\n",
    "        # Write cylinders detected by the scanner.\n",
    "        write_cylinders(f, \"D C\", [(obs[2][0], obs[2][1])\n",
    "                                   for obs in observations])\n",
    "\n",
    "    f.close()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}