def solve_lqr(A, B, Q, R, X0=np.array([0., 0.]), num_time_steps=500):
'''Solves a discrete, finite horizon LQR problem given system dynamics in
state space representation.
Arguments:
A, B: nxn state transition matrix and nxp control input matrix
Q, R: nxn state cost matrix and pxp control cost matrix
X0: initial state (n-vector)
num_time_steps: number of time steps, T
Returns:
x_sol, u_sol: Txn array of states and Txp array of controls
'''
n = np.size(A, 0)
p = np.size(B, 1)
# noise models
prior_noise = gtsam.noiseModel_Constrained.All(n)
dynamics_noise = gtsam.noiseModel_Constrained.All(n)
q_noise = gtsam.dynamic_cast_noiseModel_Diagonal_noiseModel_Gaussian(
gtsam.noiseModel_Gaussian.Information(Q))
r_noise = gtsam.dynamic_cast_noiseModel_Diagonal_noiseModel_Gaussian(
gtsam.noiseModel_Gaussian.Information(R))
# note: GTSAM 4.0.2 python wrapper doesn't have 'Information'
# wrapper, use this instead if you are not on develop branch:
# `gtsam.noiseModel_Gaussian.SqrtInformation(np.sqrt(Q)))`
# Create an empty Gaussian factor graph
graph = gtsam.GaussianFactorGraph()
# Create the keys corresponding to unknown variables in the factor graph
X = []
U = []
for i in range(num_time_steps):
X.append(gtsam.symbol(ord('x'), i))
U.append(gtsam.symbol(ord('u'), i))
# set initial state as prior
graph.add(X[0], np.eye(n), X0, prior_noise)
# Add dynamics constraint as ternary factor
# A.x1 + B.u1 - I.x2 = 0
for i in range(num_time_steps - 1):
graph.add(X[i], A, U[i], B, X[i + 1], -np.eye(n), np.zeros((n)),
dynamics_noise)
# Add cost functions as unary factors
for x in X:
graph.add(x, np.eye(n), np.zeros(n), q_noise)
for u in U:
graph.add(u, np.eye(p), np.zeros(p), r_noise)
# Solve
result = graph.optimize()
x_sol = np.zeros((num_time_steps, n))
u_sol = np.zeros((num_time_steps, p))
for i in range(num_time_steps):
x_sol[i, :] = result.at(X[i])
u_sol[i] = result.at(U[i])
return x_sol, u_sol
def create_lti_fg(A, B, X0=np.array([]), u=np.array([]), num_time_steps=500):
'''Creates a factor graph with system dynamics constraints in state-space
representation:
x_{t+1} = Ax_t + Bu_t
Arguments:
A: nxn State matrix
B: nxp Input matrix
X0: initial state (n-vector)
u: Txp control inputs (optional, if not specified, no control input will
be used)
num_time_steps: number of time steps, T, to run the system
Returns:
graph, X, U: A factor graph and lists of keys X & U for the states and
control inputs respectively
'''
# Create noise models
prior_noise = gtsam.noiseModel_Constrained.All(np.size(A, 0))
dynamics_noise = gtsam.noiseModel_Constrained.All(np.size(A, 0))
control_noise = gtsam.noiseModel_Constrained.All(1)
# Create an empty Gaussian factor graph
graph = gtsam.GaussianFactorGraph()
# Create the keys corresponding to unknown variables in the factor graph
X = []
U = []
for i in range(num_time_steps):
X.append(gtsam.symbol(ord('x'), i))
U.append(gtsam.symbol(ord('u'), i))
# set initial state as prior
if X0.size > 0:
if X0.size != np.size(A, 0):
raise ValueError("X0 dim does not match state dim")
graph.add(X[0], np.eye(X0.size), X0, prior_noise)
# Add dynamics constraint as ternary factor
# A.x1 + B.u1 - I.x2 = 0
for i in range(num_time_steps - 1):
graph.add(X[i], A, U[i], B, X[i + 1], -np.eye(np.size(A, 0)),
np.zeros((np.size(A, 0))), dynamics_noise)
# Add control inputs
for i in range(len(u)):
if np.shape(u) != (num_time_steps, np.size(B, 1)):
raise ValueError("control input is wrong size")
graph.add(U[i], np.eye(np.size(B, 1)), u[i, :], control_noise)
return graph, X, U
def create_graph():
"""Create a basic linear factor graph for testing"""
graph = gtsam.GaussianFactorGraph()
x0 = X(0)
x1 = X(1)
x2 = X(2)
BETWEEN_NOISE = gtsam.noiseModel.Diagonal.Sigmas(np.ones(1))
PRIOR_NOISE = gtsam.noiseModel.Diagonal.Sigmas(np.ones(1))
graph.add(x1, np.eye(1), x0, -np.eye(1), np.ones(1), BETWEEN_NOISE)
graph.add(x2, np.eye(1), x1, -np.eye(1), 2 * np.ones(1), BETWEEN_NOISE)
graph.add(x0, np.eye(1), np.zeros(1), PRIOR_NOISE)
return graph, (x0, x1, x2)
def estimate_poses_given_rot(factors: gtsam.BetweenFactorPose3s,
rotations: gtsam.Values,
d: int = 3):
""" Estimate Poses from measurements, given rotations. From SfmProblem in shonan.
Arguments:
factors -- data structure with many BetweenFactorPose3 factors
rotations {Values} -- Estimated rotations
Returns:
Values -- Estimated Poses
"""
I_d = np.eye(d)
def R(j):
return rotations.atRot3(j) if d == 3 else rotations.atRot2(j)
def pose(R, t):
return gtsam.Pose3(R, t) if d == 3 else gtsam.Pose2(R, t)
graph = gtsam.GaussianFactorGraph()
model = gtsam.noiseModel.Unit.Create(d)
# Add a factor anchoring t_0
graph.add(0, I_d, np.zeros((d,)), model)
# Add a factor saying t_j - t_i = Ri*t_ij for all edges (i,j)
for factor in factors:
keys = factor.keys()
i, j, Tij = keys[0], keys[1], factor.measured()
measured = R(i).rotate(Tij.translation())
graph.add(j, I_d, i, -I_d, measured, model)
# Solve linear system
translations = graph.optimize()
# Convert to Values.
result = gtsam.Values()
for j in range(rotations.size()):
tj = translations.at(j)
result.insert(j, pose(R(j), tj))
return result
def get_return_cost(graph, key):
'''Returns the value function matrix at variable `key` given a graph which
goes up and including `key`, but no further (i.e. all time steps after
`key` have already been eliminated). Does so by aggregating all unary
factors on `key`. If value function is x^TPx, then this returns P.
"Return Cost" aka "Cost-to-go" aka "Value Function".
Arguments:
graph: factor graph in LTI form
key: key in the factor graph for which we want to obtain the return cost
Returns:
return_cost: return cost, an nxn array where `n` is dimension of `key`
'''
new_fg = gtsam.GaussianFactorGraph()
for i in range(graph.size()): # loop through all factors
f = graph.at(i)
if (len(f.keys()) == 1) and (f.keys()[0] == key): # collect unary factors on `key`
new_fg.push_back(f)
sol_end = new_fg.eliminateSequential()
return sol_end.back().information()
