iekf.py

import torch
from src.utils import pdump, pload, bmmt, axat, pltt, plts
import matplotlib.pyplot as plt
import os
import numpy as np
from termcolor import cprint
from src.lie_algebra import SO3


class BaseIEKF:
    """IEKF where variable are in CPU, in double, and NOT in batch.
    Record only position."""

    # gravity vector
    g = torch.Tensor([0, 0, -9.80665]).double()

    # identity matrices
    Id3 = g.new_ones(3).diag()
    IdP = g.new_ones(15).diag()

    # Jacobians
    F = g.new_zeros(15, 15)
    G = g.new_zeros(15, 12)
    G[9:12, 6:9] = Id3
    G[12:15, 9:12] = Id3
    H = g.new_zeros(9, 15)

    def __init__(self, th_max_zupt, th_min_zupt, max_omega_norm,
        max_omega, max_acc_norm, max_acc, zupt_omega_std, zupt_acc_std,
        N_init, N_normalize):

        # initial covariance
        self.P0, self.Q = self.init_cov()

        # each ZUPT outsides this threshold set ZUPT == 1
        self.th_max_zupt = th_max_zupt
        # each ZUPT outsides this threshold set ZUPT == 0
        self.th_min_zupt = th_min_zupt

        # standard deviation on ZUPT IMU measurements
        self.zupt_omega_std = zupt_omega_std
        self.zupt_acc_std = zupt_acc_std

        self.zupt_omega_cov = (self.zupt_omega_std**2)*self.g.new_ones(3)
        self.zupt_acc_cov = (self.zupt_acc_std**2)*self.g.new_ones(3)

        # each ZUPT IMU measurement outsides a threshold discards measumrent
        self.max_omega_norm = max_omega_norm
        self.max_acc_norm = max_acc_norm
        self.max_omega = max_omega
        self.max_acc = max_acc

        # number of increments for initialization
        self.N_init = N_init
        # frequency of rotation normalization
        self.N_normalize = N_normalize

        # skew symmetric matrix
        self.Wg = self.SO3.wedge(self.g)

    def forward(self, ts, us, zupts, covs):
        """Kalman filter loop"""
        dts = ts[1:] - ts[:-1]
        zupts = self.init(ts, us, zupts)
        for i in range(1, self.N):
            self.propagate(i, dts[i-1], us[i], zupts[i])
            self.update(i, us[i], covs[i], zupts[i])
            if i % self.N_normalize == 0:
                self.normalize_rotation_matrix(i)

    def init_cov(self):
        """Initialize P0 and Q"""
        Q = self.g.new_ones(12).diag()
        omega_std = Q.new([2e-3, 2e-3, 2e-3])
        acc_std = Q.new([7e-2, 7e-2, 7e-2])
        b_omega_std = Q.new([5e-5, 5e-5, 5e-5])
        b_acc_std = Q.new([4e-5, 4e-5, 4e-5])

        Q[:3, :3] *= omega_std**2
        Q[3:6, 3:6] *= acc_std**2
        Q[6:9, 6:9] *= b_omega_std**2
        Q[9:12, 9:12] *= b_acc_std**2

        Rot0_std = 1e-0
        b_omega0_std = 1e-3
        b_acc0_std = 2e-2
        P0 = Q.new_zeros(15, 15)
        tmp = Q.new_ones(3).diag()
        P0[:2, :2] = Rot0_std**2 * Q.new_ones(2).diag()
        P0[9:12, 9:12] = b_omega0_std**2*tmp
        P0[12:15, 12:15] = b_acc0_std**2*tmp
        return P0, Q

    def init(self, ts, us, zupts):
        zupts = self.init_zupt(zupts)
        self.init_bias(ts, us, zupts)
        self.init_trajectory(ts, us)
        return zupts

    def init_zupt(self, zupts):
        """Do thresholding on ZUPT"""
        zupts[zupts > self.th_max_zupt] = 1
        zupts[zupts <= self.th_min_zupt] = 0
        zupts[:self.N_init] = 1 # start by stop
        return zupts

    def init_trajectory(self, ts, us):
        self.N = us.shape[0] # trajectory length
        # init trajectory to zero
        self.v = us.new_zeros(3)
        self.ps = us.new_zeros(self.N, 3)

    def init_bias(self, ts, us, zupts):
        """Init IMU bias and orientation with first measurements"""
        if zupts[:self.N_init].sum() == self.N_init:
            N_init = self.N_init
        else:
            N_init = torch.where(zupts[:self.N_init] == 0)[0][0].item()
            cprint('Bias initialized with only {} samples'.format(N_init),
                'yellow')
        u = us[:N_init].mean(dim=0)
        gravity = -u[3:6]
        self.Rot = self.SO3.from_2vectors(gravity, self.g)
        self.b_omega = u[:3]
        self.b_acc = self.Rot.t().mv(self.g) - gravity

        self.P = self.P0.clone()
        # init covariance
        H = u.new_zeros(6, 15)
        H[:3, 9:12] = self.Id3
        H[3:6, :3] = self.Rot.t().mm(self.Wg)
        H[3:6, 12:15] = -self.Id3
        R = self.Q[:6, :6]
        for i in range(N_init):
            S = axat(H, self.P) + R
            Kt, _ = torch.solve(self.P.mm(H.t()).t(), S)
            K = Kt.t()
            I_KH = self.IdP - K.mm(H)
            self.P = axat(I_KH, self.P.clone()) + axat(K, R)
        self.P = (self.P + self.P.t()).clone()/2

    def propagate(self, i, dt, u, zupt):
        self.propagate_cov(i, dt, u, zupt)
        z = 1 - zupt
        acc = z*(self.Rot.mv(u[3:6] - self.b_acc) + self.g)
        self.Rot = self.Rot.mm(self.SO3.exp(z*(u[:3] - self.b_omega)*dt))
        self.ps[i] = self.ps[i-1] + z*self.v*dt + 1/2*acc*(dt**2)
        self.v = self.v + acc*dt

    def propagate_cov(self, i, dt, u, zupt):
        F = self.F
        G = self.G
        Rot = self.Rot

        z = 1 - zupt
        F[3:6, :3] = z*self.Wg
        F[3:6, 12:15] = -z*Rot
        F[6:9, 3:6] = z*self.Id3
        G[3:6, 3:6] = z*Rot

        v_skew_rot = self.SO3.wedge(self.v).mm(Rot)
        p_skew_rot = self.SO3.wedge(self.ps[i-1]).mm(Rot)
        tmp = z*torch.cat((Rot, v_skew_rot, p_skew_rot))
        G[:9, :3] = tmp
        F[:9, 9:12] = -tmp

        Phi = self.IdP + F*dt + 1/2*F.mm(F)*(dt**2)
        P = axat(Phi, self.P + axat(G*dt, self.Q))
        self.P = (P + P.t())/2

    def update(self, i, u, cov, zupt):
        H = self.H
        if zupt == 1:
            z = 1
        else:
            z = 0

        H[:3, 3:6] = self.Rot.t()
        H[0, 3:6] *= z

        self.r = torch.cat((- self.Rot.t().mv(self.v),
            u[:3] - self.b_omega,
            u[3:6] - self.b_acc + self.Rot.t().mv(self.g)))
        self.r[0] *= z

        z *= self.r[3:6].norm() < self.max_omega_norm
        z *= self.r[3:6].abs().max() < self.max_omega
        z *= self.r[6:9].norm() < self.max_acc_norm
        z *= self.r[6:9].abs().max() < self.max_acc
        self.r[3:9] *= z

        H[3:6, 9:12] = z*self.Id3
        H[6:9, 12:15] = z*self.Id3
        H[6:9, :3] = -z*self.Rot.t().mm(self.Wg)

        R = torch.diag(torch.cat((cov,
            self.zupt_omega_cov,
            self.zupt_acc_cov), 0))
        S = axat(H, self.P) + R
        Kt, _ = torch.solve(self.P.mm(H.t()).t(), S)
        K = Kt.t()
        self.xi = K.mv(self.r)
        self.state_update(i)
        self.covariance_update(i, K, H, R)

    def state_update(self, i):
        Rot, Xi = self.SE3.exp2(self.xi[:9])
        self.Rot = Rot.mm(self.Rot)
        self.v = Rot.mv(self.v) + Xi[:, 0]
        self.ps[i] = Rot.mv(self.ps[i].clone()) + Xi[:, 1]
        self.b_omega += self.xi[9:12]
        self.b_acc += self.xi[12:15]

    def covariance_update(self, i, K, H, R):
        I_KH = self.IdP - K.mm(H)
        P = axat(I_KH, self.P) + axat(K, R)
        self.P = (P + P.t())/2

    def normalize_rotation_matrix(self, i):
        U, _, V = torch.svd(self.Rot)
        S = self.Id3.clone()
        S[2, 2] = torch.det(U) * torch.det(V)
        self.Rot = U.mm(S).mm(V.t())

    def nets2iekf(self, net, us, Nshift, zupts):
        """Nets input to KF input"""
        us = us.unsqueeze(0).float().cuda()

        net = net.cuda().float()
        with torch.no_grad():
            covs = net(us).squeeze()
        net = net.cpu()

        where_zupts = (zupts >= self.th_max_zupt).double()
        tmp = zupts.new_zeros(zupts.shape[0])
        tmp[499:] = torch.nn.functional.conv1d(
            where_zupts.t().unsqueeze(0),
            zupts.new_ones(1, 1, 500).double()).squeeze()
        tmp = tmp >= 500

        us = us.squeeze()
        us = us[Nshift:].double().cpu()
        covs = covs[Nshift:].double().cpu()
        zupts = zupts[Nshift:].double().cpu()
        return us, zupts, covs

    class SO3:
        # tolerance criterion
        TOL = 1e-8
        Id = torch.eye(3)

        @staticmethod
        def outer(a, b):
            return torch.einsum('i, j -> ij', a, b)

        @classmethod
        def exp(cls, phi):
            angle = phi.norm()
            if angle < cls.TOL:
                return cls.Id + cls.wedge(phi)

            axis = phi / angle
            c = angle.cos()
            s = angle.sin()
            Rot = c*cls.Id + (1-c)*cls.outer(axis, axis) + s*cls.wedge(axis)
            return Rot

        @classmethod
        def inv_left_jacobian(cls, phi):
            angle = phi.norm()
            if angle < cls.TOL:
                return cls.Id - 1/2 * cls.wedge(phi)

            axis = phi/angle
            half_angle = angle/2
            cot = 1/half_angle.tan()

            J = half_angle*cot*cls.Id + (1-half_angle*cot)*cls.outer(axis, axis) - \
                half_angle*cls.wedge(axis)
            return J

        @classmethod
        def left_jacobian(cls, phi):
            angle = phi.norm()
            if angle < cls.TOL:
                return cls.Id + 1/2*cls.wedge(phi)

            axis = phi / angle
            s = angle.sin()
            c = angle.cos()

            J = (s/angle)*cls.Id + (1-s/angle)*cls.outer(axis, axis) +\
                ((1-c)/angle) * cls.wedge(axis)
            return J

        @classmethod
        def log(cls, Rot):
            cos_angle = (0.5 * Rot.trace() - 0.5).clamp(-1., 1.)
            # Clip cos(angle) to its proper domain to avoid NaNs from rounding
            # errors
            angle = cos_angle.acos()
            if angle < cls.TOL:
                return cls.vee(Rot - cls.Id)

            angle = angle.unsqueeze(0).unsqueeze(1)
            return cls.vee((0.5 * angle/angle.sin())*(Rot-Rot.t()))

        @staticmethod
        def vee(Phi):
            return torch.stack((Phi[2, 1],
                                Phi[0, 2],
                                Phi[1, 0]), dim=0)

        @staticmethod
        def wedge(phi):
            return phi.new([[0., -phi[2], phi[1]],
                            [phi[2], 0., -phi[0]],
                            [-phi[1], phi[0], 0.]])

        @classmethod
        def from_rpy(cls, roll, pitch, yaw):
            return cls.rotz(yaw).mm(cls.roty(pitch).mm(cls.rotx(roll)))

        @classmethod
        def rotx(cls, angle_in_radians):
            c = angle_in_radians.cos()
            s = angle_in_radians.sin()
            mat = c.new_zeros((3, 3))
            mat[0, 0] = 1
            mat[1, 1] = c
            mat[2, 2] = c
            mat[1, 2] = -s
            mat[2, 1] = s
            return mat

        @classmethod
        def roty(cls, angle_in_radians):
            c = angle_in_radians.cos()
            s = angle_in_radians.sin()
            mat = c.new_zeros((3, 3))
            mat[1, 1] = 1
            mat[0, 0] = c
            mat[2, 2] = c
            mat[0, 2] = s
            mat[2, 0] = -s
            return mat

        @classmethod
        def rotz(cls, angle_in_radians):
            c = angle_in_radians.cos()
            s = angle_in_radians.sin()
            mat = c.new_zeros((3, 3))
            mat[2, 2] = 1
            mat[0, 0] = c
            mat[1, 1] = c
            mat[0, 1] = -s
            mat[1, 0] = s
            return mat

        @classmethod
        def isclose(cls, x, y):
            return (x-y).abs() < cls.TOL

        @classmethod
        def to_rpy(cls, Rot):
            """Convert a rotation matrix to RPY Euler angles."""

            pitch = torch.atan2(-Rot[2, 0],
                                torch.sqrt(Rot[0, 0]**2 + Rot[1, 0]**2))
            near_pi_over_two_mask = cls.isclose(pitch, np.pi / 2.)
            near_neg_pi_over_two_mask = cls.isclose(pitch, -np.pi / 2.)

            if near_pi_over_two_mask:
                yaw = 0.
                roll = torch.atan2(Rot[0, 1], Rot[1, 1])
            else:
                if near_neg_pi_over_two_mask:
                    yaw = 0.
                    roll = -torch.atan2(Rot[0, 1], Rot[1, 1])
                else:
                    sec_pitch = 1. / pitch.cos()
                    yaw = torch.atan2(Rot[1, 0] * sec_pitch,
                                      Rot[0, 0] * sec_pitch)
                    roll = torch.atan2(Rot[2, 1] * sec_pitch,
                                       Rot[2, 2] * sec_pitch)
                    rpys = torch.stack([roll, pitch, yaw])

            return rpys

        @classmethod
        def from_2vectors(cls, v1, v2):
            """ Returns a Rotation matrix between vectors 'v1' and 'v2'    """
            v1 = v1/v1.norm()
            v2 = v2/v2.norm()
            v = torch.cross(v1, v2)
            cosang = (v1*v2).sum()
            sinang = v.norm()
            W = cls.wedge(v)
            Rot = cls.Id + W + (1-cosang)/(sinang**2)*W.mm(W)
            return Rot

        @classmethod
        def from_quaternion(cls, quat, ordering='wxyz'):
            """Form a rotation matrix from a unit length quaternion.
            Valid orderings are 'xyzw' and 'wxyz'.
            """

            if ordering is 'xyzw':
                qx = quat[0]
                qy = quat[1]
                qz = quat[2]
                qw = quat[3]
            elif ordering is 'wxyz':
                qw = quat[0]
                qx = quat[1]
                qy = quat[2]
                qz = quat[3]

            # Form the matrix
            mat = quat.new_empty(3, 3)

            qx2 = qx * qx
            qy2 = qy * qy
            qz2 = qz * qz

            mat[0, 0] = 1. - 2. * (qy2 + qz2)
            mat[0, 1] = 2. * (qx * qy - qw * qz)
            mat[0, 2] = 2. * (qw * qy + qx * qz)

            mat[1, 0] = 2. * (qw * qz + qx * qy)
            mat[1, 1] = 1. - 2. * (qx2 + qz2)
            mat[1, 2] = 2. * (qy * qz - qw * qx)

            mat[2, 0] = 2. * (qx * qz - qw * qy)
            mat[2, 1] = 2. * (qw * qx + qy * qz)
            mat[2, 2] = 1. - 2. * (qx2 + qy2)
            return mat


    class SE3:
        Id4 = torch.eye(4)
        Id5 = torch.eye(5)
        ones = Id4.new([0, 0, 0, 1])

        @classmethod
        def exp(cls, xi):
            p = BaseIEKF.SO3.left_jacobian(xi[:3]).mv(xi[3:])
            Rot = BaseIEKF.SO3.exp(xi[:3])
            return Rot, p

        @classmethod
        def exp2(cls, xi):
            Rot = BaseIEKF.SO3.exp(xi[:3])
            V = torch.stack((xi[3:6], xi[6:9]), 1)
            Xi = BaseIEKF.SO3.left_jacobian(xi[:3]).mm(V)
            return Rot, Xi

        @classmethod
        def inv(cls, chi):
            p_inv = -chi[:3, :3].t().mv(chi[:3, 3]).unsqueeze(1)
            Rot_p_inv = torch.cat((chi[:3, :3].t(), p_inv), 1)
            return torch.cat((Rot_p_inv, cls.ones), 0)

        @classmethod
        def log(cls, chi):
            """Logarithm for SE(3)"""
            phi = BaseIEKF.SO3.log(chi[:3, :3])
            rho = BaseIEKF.SO3.inv_left_jacobian(phi).mv(chi[:3, 3])
            xi = torch.cat((phi, rho))
            return xi

        @classmethod
        def log2(cls, chi):
            """Logarithm for SE_2(3)"""
            phi = BaseIEKF.SO3.log(chi[:3, :3])
            tmp = BaseIEKF.SO3.inv_left_jacobian(phi).mm(chi[:3, 3:])
            xi = torch.cat((phi, tmp[:, 0], tmp[:, 1]), 0)
            return xi


class RecorderIEKF(BaseIEKF):

    def __init__(self, **params):
        super().__init__(**params)

    def init_trajectory(self, ts, us):
        super().init_trajectory(ts, us)
        self.Rots = us.new_zeros((self.N, 3, 3))
        self.vs = us.new_zeros(self.N, 3)
        self.b_omegas = us.new_zeros(self.N, 3)
        self.b_accs = us.new_zeros(self.N, 3)

        self.rs = us.new_zeros(self.N, 9)  # residual
        self.xis = us.new_zeros(self.N, 9)  # innovation
        self.Ps = us.new_zeros(self.N, 15, 15)

        self.Rots[0] = self.Rot
        self.b_omegas[0] = self.b_omega
        self.b_accs[0] = self.b_acc
        self.Ps[0] = self.P

    def covariance_update(self, i, K, H, R):
        super().covariance_update(i, K, H, R)
        self.save(i)

    def save(self, i):
        self.Rots[i] = self.Rot
        self.vs[i] = self.v
        self.b_omegas[i] = self.b_omega
        self.b_accs[i] = self.b_acc
        self.rs[i] = self.r[:9]
        self.xis[i] = self.xi[:9]
        self.Ps[i] = self.P

    def dump(self, address, seq, zupts, covs):
        # turn cov
        J = torch.eye(9).repeat(self.Ps.shape[0], 1, 1)
        J[:, 3:6, :3] = SO3.wedge(self.vs)
        J[:, 6:9, :3] = SO3.wedge(self.ps)
        #self.Ps = axat(J, self.Ps[:, :9, :9])
        path = os.path.join(address, seq, 'iekf.p')
        mondict = {
            'Rots': self.Rots,
            'vs': self.vs,
            'ps': self.ps,
            'b_omegas': self.b_omegas,
            'b_accs': self.b_accs,
            'rs': self.rs,
            'Ps': self.Ps.diagonal(dim1=1, dim2=2),
            'zupts': zupts,
            'covs': covs,
        }
        for k, v in mondict.items():
            mondict[k] = v.float().detach().cpu()
        pdump(mondict, path)