Fix bug in lmpc_test.py and update README.md (#7)

This PR is for CDC code collation. The main changes are:

* Fix bugs in the lmpc_test.py.
* Add explanation for argument save-trajectory

README.md

@@ -46,6 +46,7 @@ This allows to test the nonlinear lmpc controller. The argparse arguments are li
 | `num_ss_points` | int | any number that is greater than `1` | history states used for learning |
 | `ss_option` | string | `space`, `time` or `all` | criteria for history states selection |
 |   `plotting`   | action |               `store_true`                |                    save plotting if true                     |
+|   `save-trajectory`   | action |               `store_true`                |                    save simulator will store the history states and inputs if true                     |
 
 
 #### Iterative lqr for iterative tasks
@@ -62,6 +63,7 @@ This allows to test the iterative ilqr controller. The argparse arguments are li
 | `num_ss_iters` | int | any number that is greater than `1` | iterations used for learning |
 | `num_ss_points` | int | any number that is greater than `1` | history states used for learning |
 |   `plotting`   | action |               `store_true`                |                    save plotting if true                     |
+|   `save-trajectory`   | action |               `store_true`                |                    save simulator will store the history states and inputs if true                     |
 
 #### Known Issues
 To change the simulation timestep, the number of prediction horizon and number of history states used for learning should be changed.

iterative_ilqr/tests/nlmpc_test.py

@@ -1,3 +1,4 @@
+import pickle as pkl
 import numpy as np
 from utils import base
 from utils.constants_kinetic_bicycle import *
@@ -25,7 +26,7 @@ def nlmpc_test(args):
             ss_optioin = "spaceVarying"
         elif args["ss_option"] == "time":
             ss_optioin = "timeVarying"
-    x0 = [0, 0, 0, 0]
+    x0 = np.zeros((X_DIM,))
     ego = base.KineticBicycle(system_param=base.KineticBicycleParam())
     ego.set_state(x0)
     ego.set_timestep(dt)
@@ -47,9 +48,7 @@ def nlmpc_test(args):
     )
     lmpc = base.LMPC(lmpc_param, obstacle, system_param=base.KineticBicycleParam())
     lmpc.add_trajectory(ego.xcl, ego.ucl)
-    lmpc.set_initial_traj(ego.xcl, ego.ucl)
     lmpc.set_timestep(dt)
-    lmpc.set_state(x0)
     ego.set_ctrl_policy(lmpc)
     simulator = base.Simulator()
     simulator.set_robotic(ego)
@@ -58,25 +57,27 @@ def nlmpc_test(args):
     for iter in range(lap_number):
         print("iteration ", iter, "begins")
         simulator.sim(iter, sim_time=sim_time)
-        lmpc.add_trajectory(np.array(ego.all_xs[-1]).T, np.array(ego.all_inputs[-1]).T)
+        lmpc.add_trajectory(ego.data["state"][-1], ego.data["input"][-1])
     print("time at iteration 0 is", len(ego.xcl.T) * dt, " s")
     if save_lmpc_traj == True:
         np.savetxt(
-            "data/nlmpc_closed_loop_multi_laps.txt",
+            "data/lmpc_closed_loop_multi_laps.txt",
             np.round(np.array(ego.data["state"][-1]), decimals=5),
             fmt="%f",
         )
         np.savetxt(
-            "data/nlmpc_input_multi_laps.txt",
+            "data/lmpc_input_multi_laps.txt",
             np.round(np.array(ego.data["input"][-1]), decimals=5),
             fmt="%f",
         )
-    for id in range(len(ego.all_times)):
+    for id in range(len(ego.data["timestamp"])):
         lap = id + 1
-        print("time at iteration ", lap, " is ", (len(ego.all_times[id]) * dt), " s")
+        print("time at iteration ", lap, " is ", (len(ego.data["timestamp"][id]) * dt), " s")
     if args["plotting"]:
         simulator.plot_inputs()
         simulator.plot_simulation()
+    with open("data/ego_nlmpc_test.obj", "wb") as handle:
+        pkl.dump(ego, handle, protocol=pkl.HIGHEST_PROTOCOL)
 
 
 if __name__ == "__main__":
@@ -88,5 +89,6 @@ def nlmpc_test(args):
     parser.add_argument("--num-ss-iters", type=int)
     parser.add_argument("--ss-option", type=str)
     parser.add_argument("--plotting", action="store_true")
+    parser.add_argument("--save-trajectory", action="store_true")
     args = vars(parser.parse_args())
     nlmpc_test(args)