support predicted orientation

code/sim/utils/sim_utils.py

@@ -57,10 +57,13 @@ def traj2control(plan_traj, info):
     """
     plan_traj_stats = np.zeros((plan_traj.shape[0]+1, 5))
     plan_traj_stats[1:, :2] = plan_traj[:, [1,0]]
-    prev_a, prev_b = 0, 0
-    for i, (a, b) in enumerate(plan_traj):
-        rot = np.arctan((b - prev_b)/(a - prev_a))
-        plan_traj_stats[i+1, 2] = rot
+    if plan_traj.shape[1] == 4:
+        plan_traj_stats[1:, 2] = plan_traj[:, 2]
+    else:
+        prev_a, prev_b = 0, 0
+        for i, (a, b) in enumerate(plan_traj):
+            rot = np.arctan((b - prev_b)/(a - prev_a))
+            plan_traj_stats[i+1, 2] = rot
     curr_stat = np.array(
         [0, 0, 0, info['ego_velo'], info['ego_steer']]
     )