def test_lqr_linear_unbounded():
npr.seed(1)
n_batch = 2
n_state, n_ctrl = 3, 4
n_sc = n_state + n_ctrl
T = 5
C = npr.randn(T, n_batch, n_sc, n_sc)
C = np.matmul(C.transpose(0, 1, 3, 2), C)
c = npr.randn(T, n_batch, n_sc)
alpha = 0.2
R = np.tile(np.eye(n_state)+alpha*np.random.randn(n_state, n_state),
(T, n_batch, 1, 1))
S = np.tile(np.random.randn(n_state, n_ctrl), (T, n_batch, 1, 1))
F = np.concatenate((R, S), axis=3)
f = np.tile(npr.randn(n_state), (T, n_batch, 1))
x_init = npr.randn(n_batch, n_state)
# u_lower = -100.*npr.random((T, n_batch, n_ctrl))
# u_upper = 100.*npr.random((T, n_batch, n_ctrl))
u_lower = -1e4*np.ones((T, n_batch, n_ctrl))
u_upper = 1e4*np.ones((T, n_batch, n_ctrl))
tau_cp, objs_cp = lqr_cp(
C[:,0], c[:,0], F[:,0], f[:,0], x_init[0], T, n_state, n_ctrl,
None, None
)
tau_cp = tau_cp.T
x_cp = tau_cp[:,:n_state]
u_cp = tau_cp[:,n_state:]
C, c, R, S, F, f, x_init, u_lower, u_upper = [
Variable(torch.Tensor(x).double()) if x is not None else None
for x in [C, c, R, S, F, f, x_init, u_lower, u_upper]
]
dynamics = AffineDynamics(R[0,0], S[0,0], f[0,0])
u_lqr = None
x_lqr, u_lqr, objs_lqr = mpc.MPC(
n_state, n_ctrl, T, u_lower, u_upper, u_lqr,
lqr_iter=10,
backprop=False,
verbose=1,
exit_unconverged=True,
)(x_init, QuadCost(C, c), dynamics)
tau_lqr = torch.cat((x_lqr, u_lqr), 2)
tau_lqr = util.get_data_maybe(tau_lqr)
npt.assert_allclose(tau_cp, tau_lqr[:,0].numpy(), rtol=1e-3)
u_lqr = None
x_lqr, u_lqr, objs_lqr = mpc.MPC(
n_state, n_ctrl, T, None, None, u_lqr,
lqr_iter=10,
backprop=False,
exit_unconverged=False,
)(x_init, QuadCost(C, c), dynamics)
tau_lqr = torch.cat((x_lqr, u_lqr), 2)
tau_lqr = util.get_data_maybe(tau_lqr)
npt.assert_allclose(tau_cp, tau_lqr[:,0].numpy(), rtol=1e-3)
def get_frame(self, state):
""" Get a frame to use in generating a video of the results """
state = util.get_data_maybe(state.view(-1))
assert len(state) == 5
# Parse the current states from the state tensor
x, dx, cos_th, sin_th, dth = torch.unbind(state)
gravity, masscart, masspole, length = torch.unbind(self.params)
th = np.arctan2(sin_th, cos_th)
th_x = sin_th * length * 2
th_y = cos_th * length * 2
fig, ax = plt.subplots(figsize=(6, 4), dpi=300)
plt.axis(
'equal'
) # This will make the distances represented on the axes equal
# Add a rectangle representing the cart
cart = mpatches.Rectangle((x - 0.2, 0 - 0.1), 0.4, 0.2, zorder=3)
ax.add_patch(cart)
# Then, plot the pole
ax.plot((x, x + th_x), (0, th_y), color='k', zorder=10)
ax.set_xlim((-3., 3.))
ax.set_ylim((-2., 2.))
return fig, ax
def forward_numpy(C, c, x_init, u_lower, u_upper, fc0b):
_C, _c, _x_init, _u_lower, _u_upper, fc0b = [
Variable(torch.Tensor(x).double(), requires_grad=True)
if x is not None else None
for x in [C, c, x_init, u_lower, u_upper, fc0b]
]
dynamics.fcs[0].bias.data[:] = fc0b.data
# dynamics.A.data[:] = fc0b.view(n_state, n_state).data
u_init = None
x_lqr, u_lqr, objs_lqr = mpc.MPC(
n_state,
n_ctrl,
T,
_u_lower,
_u_upper,
u_init,
lqr_iter=40,
verbose=-1,
exit_unconverged=True,
backprop=False,
max_linesearch_iter=1,
slew_rate_penalty=1.0,
)(_x_init, QuadCost(_C, _c), dynamics)
return util.get_data_maybe(u_lqr.view(-1)).numpy()
def test_lqr_linear_bounded_delta():
npr.seed(1)
n_batch = 2
n_state, n_ctrl, T = 3, 4, 5
n_sc = n_state + n_ctrl
C = npr.randn(T, n_batch, n_sc, n_sc)
C = np.matmul(C.transpose(0, 1, 3, 2), C)
c = npr.randn(T, n_batch, n_sc)
alpha = 0.2
R = np.tile(
np.eye(n_state) + alpha * np.random.randn(n_state, n_state),
(T, n_batch, 1, 1))
S = 0.01 * np.tile(np.random.randn(n_state, n_ctrl), (T, n_batch, 1, 1))
F = np.concatenate((R, S), axis=3)
f = np.tile(npr.randn(n_state), (T, n_batch, 1))
x_init = npr.randn(n_batch, n_state)
u_lower = -npr.random((T, n_batch, n_ctrl))
u_upper = npr.random((T, n_batch, n_ctrl))
tau_cp, objs_cp = lqr_cp(
C[:, 0],
c[:, 0],
F[:, 0],
f[:, 0],
x_init[0],
T,
n_state,
n_ctrl,
u_lower[:, 0],
u_upper[:, 0],
)
tau_cp = tau_cp.T
x_cp = tau_cp[:, :n_state]
u_cp = tau_cp[:, n_state:]
C, c, R, S, F, f, x_init, u_lower, u_upper = [
Variable(torch.Tensor(x).double()) if x is not None else None
for x in [C, c, R, S, F, f, x_init, u_lower, u_upper]
]
dynamics = AffineDynamics(R[0, 0], S[0, 0], f[0, 0])
delta_u = 0.1
x_lqr, u_lqr, objs_lqr = mpc.MPC(
n_state,
n_ctrl,
T,
x_init,
u_lower,
u_upper,
lqr_iter=1,
verbose=1,
delta_u=delta_u,
backprop=False,
exit_unconverged=False,
)(C, c, dynamics)
u_lqr = util.get_data_maybe(u_lqr)
assert torch.abs(u_lqr).max() <= delta_u
def get_frame(self, state):
state = util.get_data_maybe(state.view(-1))
assert len(state) == 5
x, dx, cos_th, sin_th, dth = torch.unbind(state)
gravity, masscart, masspole, length = torch.unbind(self.params)
th = np.arctan2(sin_th, cos_th)
th_x = sin_th * length * 2
th_y = cos_th * length * 2
fig, ax = plt.subplots(figsize=(6, 6))
ax.plot((x, x + th_x), (0, th_y), color='k')
ax.set_xlim((-5., 5.))
ax.set_ylim((-2., 2.))
return fig, ax
def get_frame(self, x):
x = util.get_data_maybe(x.view(-1))
assert len(x) == 3
g, m, l = torch.unbind(self.params)
l = l.data[0]
cos_th, sin_th, dth = torch.unbind(x)
th = np.arctan2(sin_th, cos_th)
x = sin_th * l
y = cos_th * l
fig, ax = plt.subplots(figsize=(6, 6))
ax.plot((0, x), (0, y), color='k')
ax.set_xlim((-l * 1.2, l * 1.2))
ax.set_ylim((-l * 1.2, l * 1.2))
return fig, ax
def forward_numpy(C, c, x_init, u_lower, u_upper, F):
_C, _c, _x_init, _u_lower, _u_upper, F = [
Variable(torch.Tensor(x).double()) if x is not None else None
for x in [C, c, x_init, u_lower, u_upper, F]
]
u_init = None
x_lqr, u_lqr, objs_lqr = mpc.MPC(
n_state, n_ctrl, T, _u_lower, _u_upper, u_init,
lqr_iter=40,
verbose=1,
exit_unconverged=True,
backprop=False,
max_linesearch_iter=2,
)(_x_init, QuadCost(_C, _c), LinDx(F))
return util.get_data_maybe(u_lqr.view(-1)).numpy()
def get_frame(self, x, ax=None):
x = util.get_data_maybe(x.view(-1))
assert len(x) == 3
l = self.params[2].item()
cos_th, sin_th, dth = torch.unbind(x)
th = np.arctan2(sin_th, cos_th)
x = sin_th * l
y = cos_th * l
if ax is None:
fig, ax = plt.subplots(figsize=(6, 6))
else:
fig = ax.get_figure()
ax.plot((0, x), (0, y), color='k')
ax.set_xlim((-l * 1.2, l * 1.2))
ax.set_ylim((-l * 1.2, l * 1.2))
return fig, ax
def get_frame(self, state):
""" Get a frame to use in generating a video of the results """
state = util.get_data_maybe(state.view(-1))
# Check that we got the right number of states
assert len(state) == 2
# Parse the current states from the state tensor
x, x_dot = torch.unbind(state)
fig, ax = plt.subplots(figsize=(6, 4), dpi=300)
plt.axis('equal') # This will make the distances represented on the axes equal
# Add a rectangle representing the mass
cart = mpatches.Rectangle((x - 0.2, 0 - 0.1), 0.4, 0.2, zorder=3)
ax.add_patch(cart)
ax.set_xlim((-3., 3.))
ax.set_ylim((-2., 2.))
return fig, ax