def one_step_control(dim, T):
def cost(t, x, u):
return jnp.dot(x, x)
def dynamics(t, x, u):
return x + u
return control.ControlSpec(cost, dynamics, T, dim, dim)
def control_from_lqr(lqr):
T, dim, _ = lqr.Q.shape
dot = jnp.dot
def cost(t, x, u):
return (dot(dot(lqr.Q[t], x), x) + dot(lqr.q[t], x) +
dot(dot(lqr.R[t], u), u) + dot(lqr.r[t], u) +
dot(dot(lqr.M[t], u), x))
def dynamics(t, x, u):
return dot(lqr.A[t], x) + dot(lqr.B[t], u)
return control.ControlSpec(cost, dynamics, T, dim, dim)
def testIlqrWithNonlinearProblem(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])
X, U = control.ilqr(num_iters, 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)))
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(self.assertAllClose, atol=1e-2, check_dtypes=True)
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)))
