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)))