add jerk cost

Author: nuwandavek

README.md

@@ -12,10 +12,11 @@ We'll be using driving segments from the [comma-steering-control](https://github
 bash ./download_dataset.sh
 
 # Test this works
-python tinyphysics.py --model_path ./models/tinyphysics.onnx --data_path ./data/00000.csv --do_sim_step --do_control_step --vis
+python tinyphysics.py --model_path ./models/tinyphysics.onnx --data_path ./data/00000.csv --do_sim_step --do_control_step --debug
 
 
 # Batch Metrics on lots of routes
+python tinyphysics.py --model_path ./models/tinyphysics.onnx --data_path ./data --num_segs 1000 --do_sim_step --do_control_step
 
 ```
 
@@ -27,3 +28,14 @@ This is a "simulated car" that has been trained to mimic a very simple physics m
 ## Controllers
 Your controller should implement an [update function](https://github.com/commaai/controls_challenge/blob/1a25ee200f5466cb7dc1ab0bf6b7d0c67a2481db/controllers.py#L2) that returns the `steer_action [-1, 1]`. This controller is then run in-loop, in the simulator to autoregressively predict the car's response.
 
+*Note: The `steerFiltered` column in the dataset is not relevant here. That was the steer command for a particular platform. We're using the dataset here only to get realistic driving scenarios wrt road roll, desired acceleration and car states (velocity, forward acceleration).*
+
+
+## Evaluation
+Each rollout will result in 2 costs:
+- `lat_accel_cost`: $\dfrac{\Sigma(actual\_lat\_accel - target\_lat\_accel)^2}{steps}$
+
+- `jerk_cost`: $\dfrac{\Sigma((actual\_lat\_accel_{t} - actual\_lat\_accel_{t-1}) / \Delta t)^2}{steps - 1}$
+
+
+Minimizing both costs are very important.

tinyphysics.py

@@ -6,16 +6,21 @@
 
 from collections import namedtuple
 from hashlib import md5
+from pathlib import Path
 from typing import List, Union, Tuple
+from tqdm import tqdm
 
 from controllers import BaseController, SimpleController
 
+
 ACC_G = 9.81
 SIM_START_IDX = 100
 CONTEXT_LENGTH = 20
 VOCAB_SIZE = 1024
 LATACCEL_RANGE = [-4, 4]
+STEER_RANGE = [-1, 1]
 MAX_ACC_DELTA = 0.5
+DEL_T = 0.1
 
 State = namedtuple('State', ['roll_lataccel', 'vEgo', 'aEgo'])
 
@@ -37,17 +42,19 @@ def clip(self, value: Union[float, np.ndarray]) -> Union[float, np.ndarray]:
 
 
 class TinyPhysicsModel:
-  def __init__(self, model_path: str) -> None:
+  def __init__(self, model_path: str, debug: bool) -> None:
     self.tokenizer = LataccelTokenizer()
     options = ort.SessionOptions()
     options.intra_op_num_threads = 1
     options.inter_op_num_threads = 1
     options.log_severity_level = 3
     if 'CUDAExecutionProvider' in ort.get_available_providers():
-      print("ONNX Runtime is using GPU")
+      if debug:
+        print("ONNX Runtime is using GPU")
       provider = ('CUDAExecutionProvider', {'cudnn_conv_algo_search': 'DEFAULT'})
     else:
-      print("ONNX Runtime is using CPU")
+      if debug:
+        print("ONNX Runtime is using CPU")
       provider = 'CPUExecutionProvider'
 
     with open(model_path, "rb") as f:
@@ -78,13 +85,15 @@ def get_current_lataccel(self, sim_states: List[State], actions: List[float], pa
 
 
 class TinyPhysicsSimulator:
-  def __init__(self, model_path: str, data_path: str, do_sim_step: bool, do_control_step: bool, controller: BaseController) -> None:
+  def __init__(self, model: TinyPhysicsModel, data_path: str, do_sim_step: bool, do_control_step: bool, controller: BaseController, debug: bool = False) -> None:
     self.data_path = data_path
-    self.sim_model = TinyPhysicsModel(model_path)
+    self.sim_model = model
     self.data = self.get_data(data_path)
     self.do_sim_step = do_sim_step
     self.do_control_step = do_control_step
     self.controller = controller
+    self.debug = debug
+    self.times = []
     self.reset()
 
   def reset(self) -> None:
@@ -124,6 +133,7 @@ def control_step(self, step_idx: int) -> None:
       action = self.controller.update(self.target_lataccel_history[step_idx], self.current_lataccel, self.state_history[step_idx])
     else:
       action = 0.
+    action = np.clip(action, STEER_RANGE[0], STEER_RANGE[1])
     self.action_history.append(action)
 
   def get_state_target(self, step_idx: int) -> Tuple[List, float]:
@@ -147,42 +157,64 @@ def plot_data(self, ax, lines, axis_labels, title) -> None:
     ax.set_xlabel(axis_labels[0])
     ax.set_ylabel(axis_labels[1])
 
-  def compute_score(self) -> float:
+  def compute_cost(self) -> float:
     target = np.array(self.target_lataccel_history)[SIM_START_IDX:]
     pred = np.array(self.current_lataccel_history)[SIM_START_IDX:]
-    return -np.mean((target - pred)**2)
 
-  def rollout(self, debug=True) -> None:
-    if debug:
+    lat_accel_cost = np.mean(((target - pred) / DEL_T)**2)
+    jerk_cost = np.mean(np.diff(pred)**2)
+    return lat_accel_cost, jerk_cost
+
+  def rollout(self) -> None:
+    if self.debug:
       plt.ion()
       fig, ax = plt.subplots(4, figsize=(12, 14))
 
     for _ in range(len(self.data)):
       self.step()
-      if debug and self.step_idx % 10 == 0:
+      if self.debug and self.step_idx % 10 == 0:
         print(f"Step {self.step_idx:<5}: Current lataccel: {self.current_lataccel:>6.2f}, Target lataccel: {self.target_lataccel_history[-1]:>6.2f}")
         self.plot_data(ax[0], [(self.target_lataccel_history, 'Target lataccel'), (self.current_lataccel_history, 'Current lataccel')], ['Step', 'Lateral Acceleration'], 'Lateral Acceleration')
         self.plot_data(ax[1], [(self.action_history, 'Action')], ['Step', 'Action'], 'Action')
         self.plot_data(ax[2], [(np.array(self.state_history)[:, 0], 'Roll Lateral Acceleration')], ['Step', 'Lateral Accel due to Road Roll'], 'Lateral Accel due to Road Roll')
         self.plot_data(ax[3], [(np.array(self.state_history)[:, 1], 'vEgo')], ['Step', 'vEgo'], 'vEgo')
         plt.pause(0.01)
 
-    if debug:
+    if self.debug:
       plt.ioff()
       plt.show()
-
-    return self.compute_score()
+    return self.compute_cost()
 
 
 if __name__ == "__main__":
   parser = argparse.ArgumentParser()
   parser.add_argument("--model_path", type=str, required=True)
   parser.add_argument("--data_path", type=str, required=True)
+  parser.add_argument("--num_segs", type=int, default=1000)
   parser.add_argument("--do_sim_step", action='store_true')
   parser.add_argument("--do_control_step", action='store_true')
-  parser.add_argument("--vis", action='store_true')
+  parser.add_argument("--debug", action='store_true')
   args = parser.parse_args()
 
-  sim = TinyPhysicsSimulator(args.model_path, args.data_path, args.do_sim_step, args.do_control_step, controller=SimpleController())
-  score = sim.rollout(args.vis)
-  print(f"Final score: {score:>6.4}")
+  tinyphysicsmodel = TinyPhysicsModel(args.model_path, debug=args.debug)
+
+  data_path = Path(args.data_path)
+  if data_path.is_file():
+    sim = TinyPhysicsSimulator(tinyphysicsmodel, args.data_path, args.do_sim_step, args.do_control_step, controller=SimpleController(), debug=args.debug)
+    lat_accel_cost, jerk_cost = sim.rollout()
+    print(f"\nAverage lat_accel_cost: {lat_accel_cost:>6.4}, average jerk_cost: {jerk_cost:>6.4}")
+  elif data_path.is_dir():
+    costs = []
+    files = sorted(data_path.iterdir())[:args.num_segs]
+    for data_file in tqdm(files, total=len(files)):
+      sim = TinyPhysicsSimulator(tinyphysicsmodel, str(data_file), args.do_sim_step, args.do_control_step, controller=SimpleController(), debug=args.debug)
+      cost = sim.rollout()
+      costs.append(cost)
+    costs = np.array(costs)
+    print(f"\nAverage lat_accel_cost: {np.mean(costs[:, 0]):>6.4}, average jerk_cost: {np.mean(costs[:, 1]):>6.4}")
+    plt.hist(costs[:, 0], bins=np.arange(0, 2, 0.1), label='lat_accel_cost', alpha=0.5)
+    plt.hist(costs[:, 1], bins=np.arange(0, 2, 0.1), label='jerk_cost', alpha=0.5)
+    plt.xlabel('costs')
+    plt.ylabel('Frequency')
+    plt.title('costs Distribution')
+    plt.show()