Bel_bar[:,t] = np.vstack((X_bar[:,t], cov2vec(Sigmas[:,:,t])))
As, Bs, Cs = car.linearize_belief_dynamics_trajectory(Bel_bar, U_bar, s, Q, R)
for t in xrange(T-1):
    # Overwrites Bel_bar
    #NOTE Keep making mistake in line below
    # Small mismatch occurs due to dynamics noise
    Bel_bar[:,t+1] = Cs[:,t]#As[:,:,t]*Bel_bar[:,t] + Bs[:,:,t]*np.mat(U_bar[:,t]).T + Cs[:,t]

car.draw_belief_trajectory(Bel_bar)
#plt.show()
#stop

# Apply SCP

rho_bel = 0.05
rho_u = 0.05
N_iter = 5
goal_bel = np.copy(Bel_bar[:,-1])
goal_bel[car.NX:] = 0

opt_bels, opt_ctrls, opt_vals = scp_solver_beliefs(s, Bel_bar.copy(), U_bar,\
                Q, R, rho_bel, rho_u, goal_bel, N_iter, method='shooting')
Bel_opt = np.mat(np.copy(Bel_bar))
for t in xrange(T-1):
    Bel_opt[:,t+1] = car.belief_dynamics(Bel_opt[:,t], opt_ctrls[:,t], s, Q, R)

car.draw_belief_trajectory(Bel_opt, color='yellow')

plt.show()

    )  # NOTE No obs

arm.draw_trajectory(X, mus, Sigmas)

Bel_bar = np.mat(np.zeros((arm.NB, T)))
for t in xrange(T):
    Bel_bar[:, t] = np.vstack((X[:, t], cov2vec(Sigmas[:, :, t])))

rho_bel = 0.2
rho_u = 0.3
N_iter = 1
goal_bel = np.copy(Bel_bar[:, -1])
goal_bel[arm.NX :] = 0

opt_bels, opt_ctrls, opt_vals = scp_solver_beliefs(
    s, Bel_bar.copy(), U, Q, R, rho_bel, rho_u, goal_bel, N_iter, arm.NX, method="shooting"
)


opt_mus = np.mat(np.zeros((arm.NX, T)))
opt_mus[:, 0] = X[:, 0]
opt_X = opt_mus.copy()
opt_Sigmas = np.zeros((Q.shape[0], Q.shape[1], T))
opt_Sigmas[:, :, 0] = Sigmas[:, :, 0]
opt_ctrls = np.mat(opt_ctrls)

for t in xrange(1, T):
    opt_X[:, t] = arm.dynamics(opt_X[:, t - 1], opt_ctrls[:, t - 1])
    opt_mus[:, t], opt_Sigmas[:, :, t] = ekf_update(
        arm.dynamics,
        lambda x: arm.observe(s, x=x),
                                      R, mus[:, t - 1], Sigmas[:, :, t - 1],
                                      U[:, t - 1], None)  #NOTE No obs

localizer.draw_trajectory(X, mus, Sigmas)

Bel_bar = np.mat(np.zeros((localizer.NB, T)))
for t in xrange(T):
    Bel_bar[:, t] = np.vstack((X[:, t], cov2vec(Sigmas[:, :, t])))

rho_bel = 0.1
rho_u = 0.1
N_iter = 5
goal_bel = np.copy(Bel_bar[:, -1])
goal_bel[localizer.NX:] = 0

opt_bels, opt_ctrls, opt_vals = scp_solver_beliefs(s, Bel_bar.copy(), U,\
               Q, R, rho_bel, rho_u, goal_bel, N_iter, localizer.NX, method='shooting')

opt_mus = np.mat(np.zeros((localizer.NX, T)))
opt_mus[:, 0] = X[:, 0]
opt_X = opt_mus.copy()
opt_Sigmas = np.zeros((Q.shape[0], Q.shape[1], T))
opt_Sigmas[:, :, 0] = Sigmas[:, :, 0]
opt_ctrls = np.mat(opt_ctrls)

for t in xrange(1, T):
    opt_X[:, t] = localizer.dynamics(opt_X[:, t - 1], opt_ctrls[:, t - 1])
    opt_mus[:, t], opt_Sigmas[:, :, t] = ekf_update(
        localizer.dynamics, lambda x: localizer.observe(s, x=x), Q, R,
        opt_mus[:, t - 1], opt_Sigmas[:, :, t - 1], opt_ctrls[:, t - 1], None)

localizer.draw_trajectory(opt_X,