def testMpcWithLqrProblemSpecifiedGenerally(self):
dim, T, num_iters = 2, 10, 3
p = one_step_control(dim, T)
x0 = self.rng.normal(size=dim)
solver = partial(control.ilqr, num_iters)
X, U = control.mpc_predict(solver, p, x0, jnp.zeros((T, dim)))
np.testing.assert_allclose(X[0], x0)
np.testing.assert_allclose(U[0], -x0)
np.testing.assert_allclose(X[1:], jnp.zeros((T, 2)))
np.testing.assert_allclose(U[1:], jnp.zeros((T - 1, 2)))
def testMpcWithLqrProblemSpecifiedGenerally(self): randn = np.random.RandomState(0).randn dim, T, num_iters = 2, 10, 3 p = one_step_control(dim, T) x0 = randn(dim) solver = partial(control.ilqr, num_iters) X, U = control.mpc_predict(solver, p, x0, jnp.zeros((T, dim))) self.assertAllClose(X[0], x0, check_dtypes=True) self.assertAllClose(U[0], -x0, check_dtypes=True) self.assertAllClose(X[1:], jnp.zeros((T, 2)), check_dtypes=True) self.assertAllClose(U[1:], jnp.zeros((T - 1, 2)), check_dtypes=True)
def testMpcWithLqrProblem(self):
dim, T, num_iters = 2, 10, 3
lqr = one_step_lqr(dim, T)
p = control_from_lqr(lqr)
x0 = jnp.array(self.rng.normal(size=dim))
solver = partial(control.ilqr, num_iters)
X, U = control.mpc_predict(solver, p, x0, jnp.zeros((T, dim)))
np.testing.assert_allclose(X[0], x0)
np.testing.assert_allclose(U[0], -x0)
np.testing.assert_allclose(X[1:], jnp.zeros((T, 2)))
np.testing.assert_allclose(U[1:], jnp.zeros((T - 1, 2)))
def testMpcWithLqrProblem(self):
randn = onp.random.RandomState(0).randn
dim, T, num_iters = 2, 10, 3
lqr = one_step_lqr(dim, T)
p = control_from_lqr(lqr)
x0 = randn(dim)
solver = partial(control.ilqr, num_iters)
X, U = control.mpc_predict(solver, p, x0, np.zeros((T, dim)))
self.assertAllClose(X[0], x0, check_dtypes=True)
self.assertAllClose(U[0], -x0, check_dtypes=True)
self.assertAllClose(X[1:], np.zeros((T, 2)), check_dtypes=True)
self.assertAllClose(U[1:], np.zeros((T - 1, 2)), check_dtypes=True)
def testMpcWithNonlinearProblem(self):
def cost(t, x, u):
return (x[0]**2. + 1e-3 * u[0]**2.) / (t + 1.)
def dynamics(t, x, u):
return (x**2. - u**2.) / (t + 1.)
T, num_iters, d = 10, 7, 1
p = control.ControlSpec(cost, dynamics, T, d, d)
x0 = jnp.array([0.2])
solver = partial(control.ilqr, num_iters)
X, U = control.mpc_predict(solver, p, x0, 1e-5 * jnp.ones((T, d)))
assert_close = partial(np.testing.assert_allclose, atol=1e-2)
assert_close(X[0], x0)
assert_close(U[0]**2., x0**2.)
assert_close(X[1:], jnp.zeros((T, d)))
assert_close(U[1:], jnp.zeros((T - 1, d)))
