Add unit tests for the wheeled inverted pendulum

tests/test_humanoid_one_step.py

@@ -12,8 +12,6 @@
 from qpmpc import MPCProblem, solve_mpc
 from qpmpc.solve_mpc import MPCQP
 
-gravity = 9.81  # [m] / [s]^2
-
 
 @dataclass
 class HumanoidSteppingProblem:
@@ -39,7 +37,7 @@ def setUp(self):
         )
         input_matrix = np.array([T**3 / 6.0, T**2 / 2.0, T])
         input_matrix = input_matrix.reshape((3, 1))
-        zmp_from_state = np.array([1.0, 0.0, -problem.com_height / gravity])
+        zmp_from_state = np.array([1.0, 0.0, -problem.com_height / 9.81])
         ineq_matrix = np.array([+zmp_from_state, -zmp_from_state])
         cur_max = problem.start_pos + 0.5 * problem.foot_length
         cur_min = problem.start_pos - 0.5 * problem.foot_length

tests/test_wheeled_inverted_pendulum.py

@@ -0,0 +1,46 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+#
+# Copyright 2022 Stéphane Caron
+# SPDX-License-Identifier: Apache-2.0
+
+import unittest
+
+import numpy as np
+
+from qpmpc import solve_mpc
+from qpmpc.systems import WheeledInvertedPendulum
+
+
+class TestWheeledInvertedPendulum(unittest.TestCase):
+    def setUp(self):
+        self.pendulum = WheeledInvertedPendulum()
+
+    def test_properties(self):
+        self.assertGreater(self.pendulum.horizon_duration, 0.1)
+        self.assertGreater(self.pendulum.omega, 0.1)
+
+    def test_mpc_problem(self):
+        mpc_problem = self.pendulum.build_mpc_problem(
+            terminal_cost_weight=10.0,
+            stage_state_cost_weight=1.0,
+            stage_input_cost_weight=1e-3,
+        )
+        initial_state = np.zeros(self.pendulum.STATE_DIM)
+        goal_state = initial_state.copy()
+        nx = self.pendulum.STATE_DIM
+        target_states = np.zeros(self.pendulum.nb_timesteps * nx)
+        mpc_problem.update_initial_state(initial_state)
+        mpc_problem.update_goal_state(goal_state)
+        mpc_problem.update_target_states(target_states)
+        plan = solve_mpc(mpc_problem, solver="proxqp")
+        self.assertIsNotNone(plan)
+        self.assertIsNotNone(plan.first_input)
+        dt = self.pendulum.sampling_period
+        print(f"{plan.first_input=}")
+        state = self.pendulum.integrate(initial_state, plan.first_input, dt)
+        self.assertTrue(np.allclose(state, initial_state))
+
+
+if __name__ == "__main__":
+    unittest.main()