From e59adfd62cc685b5ed44b9e4bc0f2fc62617b964 Mon Sep 17 00:00:00 2001
From: Emmanuel Akinlosotu <emmanuelakinlosotu12@gmail.com>
Date: Tue, 13 May 2025 05:01:33 -0400
Subject: [PATCH] Implement tuned PID controller with feed-forward term (score
 ~96.7)

---
 .gitignore               |  75 +++++++++++++-
 controllers/tuned_pid.py |  40 ++++++++
 requirements.txt         |   3 +
 scripts/run_tuning.bat   |  34 +++++++
 scripts/run_tuning.sh    |  34 +++++++
 scripts/tune_pid.py      | 215 +++++++++++++++++++++++++++++++++++++++
 6 files changed, 398 insertions(+), 3 deletions(-)
 create mode 100644 controllers/tuned_pid.py
 create mode 100644 scripts/run_tuning.bat
 create mode 100755 scripts/run_tuning.sh
 create mode 100755 scripts/tune_pid.py

diff --git a/.gitignore b/.gitignore
index 534a3a27..bb7627ac 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,5 +1,74 @@
+# Python
 __pycache__/
-data/
+*.py[cod]
+*$py.class
+*.so
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+
+# Virtual Environment
+venv/
+env/
+ENV/
+
+# Project specific
 report.html
-models/tinyphysics_*.onnx
-*.swp
+best_pid_params.txt
+visualizations/
+*.onnx
+
+# Jupyter Notebooks
+.ipynb_checkpoints
+*.ipynb_checkpoints/
+*.ipynb
+
+# Visual Studio Code
+.vscode/
+*.code-workspace
+.history/
+
+# PyCharm
+.idea/
+*.iml
+*.iws
+*.ipr
+
+# Logs
+logs/
+*.log
+
+# Data (optional - uncomment if you don't want to track data files)
+data/
+*.csv
+
+# Outputs
+*.png
+*.jpg
+*.jpeg
+*.gif
+*.svg
+*.pkl
+*.joblib
+
+# OS specific
+.DS_Store
+.DS_Store?
+._*
+.Spotlight-V100
+.Trashes
+ehthumbs.db
+Thumbs.db
diff --git a/controllers/tuned_pid.py b/controllers/tuned_pid.py
new file mode 100644
index 00000000..edf43880
--- /dev/null
+++ b/controllers/tuned_pid.py
@@ -0,0 +1,40 @@
+from . import BaseController
+import numpy as np
+
+class Controller(BaseController):
+  """
+  A simplified tuned PID controller with feed-forward term
+  
+  Based on the baseline PID controller but adds a feed-forward term
+  Feed-forward term: k_ff * target_lataccel / v_ego**2
+  """
+  def __init__(self):
+    # Optimized PID gains from Optuna (trial 25)
+    self.p = 0.21942091482230813  # Proportional gain
+    self.i = 0.09438349898803905  # Integral gain
+    self.d = -0.06219594603865014  # Derivative gain
+    
+    # Add feed-forward gain
+    self.k_ff = 0.02737288447587737
+    
+    # State variables
+    self.error_integral = 0
+    self.prev_error = 0
+
+  def update(self, target_lataccel, current_lataccel, state, future_plan):
+      # Extract vehicle velocity for feed-forward term
+      v_ego = state.v_ego
+      
+      # Calculate PID terms (same as baseline)
+      error = (target_lataccel - current_lataccel)
+      self.error_integral += error
+      error_diff = error - self.prev_error
+      self.prev_error = error
+      
+      # Calculate feed-forward term (avoid division by zero)
+      feed_forward = 0
+      if v_ego > 0.1:  # Only apply feed-forward when moving
+          feed_forward = self.k_ff * target_lataccel / (v_ego**2)
+      
+      # Combine PID and feed-forward
+      return self.p * error + self.i * self.error_integral + self.d * error_diff + feed_forward 
\ No newline at end of file
diff --git a/requirements.txt b/requirements.txt
index ff7e4bf4..0d9078c1 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -4,3 +4,6 @@ pandas==2.1.2
 matplotlib==3.8.1
 seaborn==0.13.2
 tqdm
+optuna==3.5.0
+plotly==5.26.0
+kaleido==0.2.1
diff --git a/scripts/run_tuning.bat b/scripts/run_tuning.bat
new file mode 100644
index 00000000..6df0c6fa
--- /dev/null
+++ b/scripts/run_tuning.bat
@@ -0,0 +1,34 @@
+@echo off
+echo Starting PID Controller Tuning Process
+echo ==================================
+
+REM Install requirements
+pip install -r requirements.txt
+
+REM Activate environment if needed
+REM call venv\Scripts\activate.bat
+
+REM Run PID tuning with recommended parameters
+python scripts/tune_pid.py ^
+  --model_path ./models/tinyphysics.onnx ^
+  --data_path ./data ^
+  --num_segs 200 ^
+  --n_trials 100 ^
+  --update_controller
+
+echo.
+echo Tuning completed! Updated controller with best parameters.
+echo Running final evaluation against baseline PID...
+echo.
+
+REM Generate final evaluation report
+python eval.py ^
+  --model_path ./models/tinyphysics.onnx ^
+  --data_path ./data ^
+  --num_segs 100 ^
+  --test_controller tuned_pid ^
+  --baseline_controller pid
+
+echo.
+echo Evaluation complete! See report.html for results.
+echo ================================== 
\ No newline at end of file
diff --git a/scripts/run_tuning.sh b/scripts/run_tuning.sh
new file mode 100755
index 00000000..a64b4841
--- /dev/null
+++ b/scripts/run_tuning.sh
@@ -0,0 +1,34 @@
+#!/bin/bash
+echo "Starting PID Controller Tuning Process"
+echo "=================================="
+
+# Install requirements
+pip install -r requirements.txt
+
+# Activate environment if needed
+# source venv/bin/activate
+
+# Run PID tuning with recommended parameters
+python scripts/tune_pid.py \
+  --model_path ./models/tinyphysics.onnx \
+  --data_path ./data \
+  --num_segs 200 \
+  --n_trials 100 \
+  --update_controller
+
+echo
+echo "Tuning completed! Updated controller with best parameters."
+echo "Running final evaluation against baseline PID..."
+echo
+
+# Generate final evaluation report
+python eval.py \
+  --model_path ./models/tinyphysics.onnx \
+  --data_path ./data \
+  --num_segs 100 \
+  --test_controller tuned_pid \
+  --baseline_controller pid
+
+echo
+echo "Evaluation complete! See report.html for results."
+echo "==================================" 
\ No newline at end of file
diff --git a/scripts/tune_pid.py b/scripts/tune_pid.py
new file mode 100755
index 00000000..2d51c411
--- /dev/null
+++ b/scripts/tune_pid.py
@@ -0,0 +1,215 @@
+#!/usr/bin/env python
+"""
+PID Controller Parameter Tuning with Optuna
+
+This script uses Optuna to find optimal gains for the tuned_pid controller
+by running the simulator on a subset of the data segments.
+
+Usage:
+    python scripts/tune_pid.py --model_path models/tinyphysics.onnx --data_path data --num_segs 200
+"""
+
+import argparse
+import importlib
+import os
+import sys
+import time
+from functools import partial
+from pathlib import Path
+
+import numpy as np
+import optuna
+from optuna.samplers import TPESampler
+from optuna.visualization import plot_optimization_history, plot_param_importances
+
+# Add parent directory to path to allow importing tinyphysics
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from tinyphysics import TinyPhysicsModel, TinyPhysicsSimulator
+
+# Import controller
+from controllers.tuned_pid import Controller
+
+def objective(trial, model_path, data_files, num_segs=200, lataccel_weight=50):
+    """Optuna objective function that evaluates PID parameters"""
+    # Sample hyperparameters - main gains
+    p = trial.suggest_float("p", 0.05, 0.8, log=True)
+    i = trial.suggest_float("i", 0.01, 0.3, log=True)
+    d = trial.suggest_float("d", -0.5, 0.0, log=False)
+    k_ff = trial.suggest_float("k_ff", 0.005, 0.2, log=True)
+    
+    # Setup the model
+    tinyphysicsmodel = TinyPhysicsModel(model_path, debug=False)
+    
+    # Shuffle and select a subset of data files
+    np.random.shuffle(data_files)
+    selected_data_files = data_files[:num_segs]
+    
+    # Collect costs for each segment
+    costs = []
+    for i_seg, data_file in enumerate(selected_data_files):
+        # Progress indicator
+        if i_seg % 10 == 0:
+            print(f"Progress: {i_seg}/{len(selected_data_files)} segments")
+        
+        # Create controller with the trial parameters
+        controller = Controller()
+        # Update controller parameters
+        controller.p = p
+        controller.i = i
+        controller.d = d
+        controller.k_ff = k_ff
+        
+        # Create simulator and run rollout
+        sim = TinyPhysicsSimulator(tinyphysicsmodel, str(data_file), controller=controller, debug=False)
+        cost = sim.rollout()
+        costs.append(cost)
+    
+    # Calculate the mean costs across all segments
+    if costs:
+        lataccel_costs = np.mean([c['lataccel_cost'] for c in costs])
+        jerk_costs = np.mean([c['jerk_cost'] for c in costs]) 
+        total_cost = np.mean([c['total_cost'] for c in costs])
+        
+        # We can weight different aspects differently during tuning
+        # tuning_cost = lataccel_weight * lataccel_costs + jerk_costs
+        tuning_cost = total_cost  # Default to using simulator's total_cost
+    else:
+        tuning_cost = float('inf')  # If no valid costs, this is a bad trial
+    
+    print(f"Trial {trial.number}: p={p:.4f}, i={i:.4f}, d={d:.4f}, k_ff={k_ff:.4f}, " +
+          f"cost={tuning_cost:.4f}")
+    
+    return tuning_cost
+
+def run_tuning(model_path, data_path, num_segs=200, n_trials=100, lataccel_weight=50):
+    """Run the Optuna tuning process"""
+    # Get data files
+    data_files = list(Path(data_path).glob('*.csv'))
+    if not data_files:
+        raise ValueError(f"No data files found in {data_path}")
+    
+    print(f"Found {len(data_files)} data files, will use up to {num_segs} for each trial")
+    
+    # Create Optuna study
+    study = optuna.create_study(
+        direction="minimize",
+        sampler=TPESampler(seed=42),
+        pruner=optuna.pruners.MedianPruner(n_warmup_steps=10)
+    )
+    
+    # Run optimization
+    objective_func = partial(objective, model_path=model_path, 
+                           data_files=data_files, num_segs=num_segs,
+                           lataccel_weight=lataccel_weight)
+                           
+    study.optimize(objective_func, n_trials=n_trials)
+    
+    # Print results
+    best_params = study.best_params
+    best_value = study.best_value
+    
+    print("\n" + "="*50)
+    print("Best Parameters:")
+    print(f"p: {best_params['p']:.6f}")
+    print(f"i: {best_params['i']:.6f}")
+    print(f"d: {best_params['d']:.6f}")
+    print(f"k_ff: {best_params['k_ff']:.6f}")
+    print(f"Best Total Cost: {best_value:.4f}")
+    print("="*50)
+    
+    # Save best parameters to a file
+    with open("best_pid_params.txt", "w") as f:
+        f.write(f"p = {best_params['p']:.6f}\n")
+        f.write(f"i = {best_params['i']:.6f}\n")
+        f.write(f"d = {best_params['d']:.6f}\n")
+        f.write(f"k_ff = {best_params['k_ff']:.6f}\n")
+        f.write(f"# Best Total Cost: {best_value:.4f}\n")
+    
+    # Try to generate visualization if matplotlib is available
+    try:
+        # Create visualizations directory
+        os.makedirs("visualizations", exist_ok=True)
+        
+        # Save parameter importance plot
+        param_importance_fig = plot_param_importances(study)
+        param_importance_fig.write_image("visualizations/param_importance.png")
+        
+        # Save optimization history plot
+        history_fig = plot_optimization_history(study)
+        history_fig.write_image("visualizations/optimization_history.png")
+        
+        print("Saved visualization plots to 'visualizations' directory")
+    except Exception as e:
+        print(f"Could not generate visualizations: {e}")
+    
+    return best_params, best_value
+
+def modify_tuned_pid_with_best_params(best_params):
+    """Update the tuned_pid.py file with the best parameters"""
+    controller_path = Path("controllers/tuned_pid.py")
+    with open(controller_path, "r") as f:
+        content = f.read()
+    
+    # Update the parameters in the __init__ method
+    # Find line with self.p assignment and update it
+    lines = content.split('\n')
+    for i, line in enumerate(lines):
+        if "self.p =" in line:
+            lines[i] = f"    self.p = {best_params['p']:.6f}"
+        elif "self.i =" in line:
+            lines[i] = f"    self.i = {best_params['i']:.6f}"
+        elif "self.d =" in line:
+            lines[i] = f"    self.d = {best_params['d']:.6f}"
+        elif "self.k_ff =" in line:
+            lines[i] = f"    self.k_ff = {best_params['k_ff']:.6f}"
+    
+    with open(controller_path, "w") as f:
+        f.write('\n'.join(lines))
+    
+    print(f"Updated {controller_path} with the best parameters")
+
+def main():
+    parser = argparse.ArgumentParser(description="Tune PID controller parameters using Optuna")
+    parser.add_argument("--model_path", required=True, help="Path to the model file")
+    parser.add_argument("--data_path", required=True, help="Path to the data directory")
+    parser.add_argument("--num_segs", type=int, default=200, help="Number of segments to use for tuning")
+    parser.add_argument("--n_trials", type=int, default=100, help="Number of Optuna trials")
+    parser.add_argument("--lataccel_weight", type=float, default=50, 
+                        help="Weight for lataccel_cost in the objective function")
+    parser.add_argument("--update_controller", action="store_true", help="Update the controller file with best params")
+    
+    args = parser.parse_args()
+    
+    start_time = time.time()
+    
+    # Make sure optuna is installed
+    try:
+        import optuna
+    except ImportError:
+        print("Optuna not found. Installing...")
+        os.system("pip install optuna")
+    
+    # Also ensure plotly is available for visualizations
+    try:
+        import plotly
+    except ImportError:
+        print("Plotly not found. Installing...")
+        os.system("pip install plotly kaleido")
+    
+    best_params, best_value = run_tuning(
+        model_path=args.model_path,
+        data_path=args.data_path,
+        num_segs=args.num_segs,
+        n_trials=args.n_trials,
+        lataccel_weight=args.lataccel_weight
+    )
+    
+    # Optionally update the controller file
+    if args.update_controller:
+        modify_tuned_pid_with_best_params(best_params)
+    
+    elapsed_time = time.time() - start_time
+    print(f"Tuning completed in {elapsed_time/60:.2f} minutes")
+
+if __name__ == "__main__":
+    main() 
\ No newline at end of file