"""HPI Control Validation — Inverted Pendulum (NN edition) Runs HPI algorithm (NN-based), compares with LQR baseline, validates in closed loop, outputs trajectories to CSV, and plots results. """ import os import numpy as np import pandas as pd import matplotlib matplotlib.use("TkAgg") import matplotlib.pyplot as plt from scipy.integrate import solve_ivp from scipy.linalg import solve_continuous_are from hpi.hpi_controller import HPIController plt.rcParams.update({ "figure.figsize": (12, 8), "font.size": 12, "axes.titlesize": 14, "axes.labelsize": 13, "legend.fontsize": 11, "figure.dpi": 120, }) OUTPUT_DIR = "output" os.makedirs(OUTPUT_DIR, exist_ok=True) # ═══════════════════════════════════════════════════════════════════ # LQR Baseline # ═══════════════════════════════════════════════════════════════════ def compute_lqr(): """Compute LQR solution for linearized inverted pendulum.""" A = np.array([[0., 1.], [1., 0.]]) B = np.array([[0.], [1.]]) Q = np.eye(2) R = np.array([[1.]]) P = solve_continuous_are(A, B, Q, R) K = (np.linalg.solve(R, B.T @ P)).ravel() return K # ═══════════════════════════════════════════════════════════════════ # HPI # ═══════════════════════════════════════════════════════════════════ def run_hpi(): """Run HPI algorithm, return controller and results.""" print("=" * 60) print(" HPI Adaptive Optimal Control — Inverted Pendulum (NN)") print("=" * 60) controller = HPIController( J=1.0, mgl=1.0, R=1.0, pe_amplitude=0.5, lr=1e-3, epochs=500, lambda_weight=1e-6, ) results = controller.run( x0=np.array([0.5, 0.0]), T=5.0, dt=0.05, L_start=0.1, L_step=1.0, L_max=10.0, gamma=0.1, epsilon=1e-3, ) print(f"\n Phase 1 L0 = {results['phase1_L0']:.2f}") print(f" Phase 1 info: {results['phase1_info']}") print(f" Phase 2 iterations: {len(results['phase2_history'])}") print(f" Phase 3 converged: {results['phase3_converged']}") print(f" Phase 3 iterations: {len(results['phase3_history'])}") return controller, results # ═══════════════════════════════════════════════════════════════════ # Simulation # ═══════════════════════════════════════════════════════════════════ def simulate_closed_loop(controller, x0, T=5.0, n_points=500): """Closed-loop simulation with HPI controller. Returns: t, X, U, V """ t_eval = np.linspace(0, T, n_points) def control(t, x): return controller.trainer.compute_u_hat(x) def ode(t, x): return controller.collector.dynamics(t, x, control) sol = solve_ivp(ode, (0, T), x0, method="RK45", t_eval=t_eval, rtol=1e-8, atol=1e-10) t = sol.t X = sol.y.T U = np.array([control(ti, xi) for ti, xi in zip(t, X)]) V = controller.trainer.compute_V(X) return t, X, U, V def simulate_lqr_closed_loop(J, mgl, K, x0, T=5.0, n_points=500): """Closed-loop simulation with LQR controller on nonlinear pendulum.""" t_eval = np.linspace(0, T, n_points) def u_lqr(t, x): return float(-np.dot(K, x)) def dynamics(t, x): return np.array([x[1], (mgl / J) * np.sin(x[0]) + (1.0 / J) * u_lqr(t, x)]) sol = solve_ivp(dynamics, (0, T), x0, method="RK45", t_eval=t_eval, rtol=1e-8, atol=1e-10) t = sol.t X = sol.y.T U = np.array([u_lqr(ti, xi) for ti, xi in zip(t, X)]) return t, X, U # ═══════════════════════════════════════════════════════════════════ # I/O # ═══════════════════════════════════════════════════════════════════ def save_csv(t, X, U, V, label): df = pd.DataFrame({"time": t, "x1": X[:, 0], "x2": X[:, 1], "u": U, "V": V}) path = os.path.join(OUTPUT_DIR, f"trajectory_{label}.csv") df.to_csv(path, index=False) return path # ═══════════════════════════════════════════════════════════════════ # Plotting # ═══════════════════════════════════════════════════════════════════ def plot_results(all_data, controller, results): """Plot state/control/value trajectories, phase portrait, convergence.""" colors = ["#1f77b4", "#d62728", "#2ca02c", "#ff7f0e", "#9467bd", "#8c564b"] # ── Figure 1: State + Control + Value ── fig1, axes1 = plt.subplots(2, 2, figsize=(14, 10)) for idx, (label, (t, X, U, V)) in enumerate(all_data.items()): c = colors[idx % len(colors)] axes1[0, 0].plot(t, X[:, 0], color=c, label=f"{label} (x0=[{X[0,0]:.1f},{X[0,1]:.1f}])") axes1[0, 1].plot(t, X[:, 1], color=c) axes1[1, 0].plot(t, U, color=c) axes1[1, 1].plot(t, V, color=c) axes1[0, 0].set_ylabel("$x_1$ (angle)") axes1[0, 0].set_title("State $x_1$") axes1[0, 0].legend(fontsize=9) axes1[0, 0].grid(True, alpha=0.3) axes1[0, 0].axhline(y=0, color="k", lw=0.5) axes1[0, 1].set_ylabel("$x_2$ (ang. velocity)") axes1[0, 1].set_title("State $x_2$") axes1[0, 1].grid(True, alpha=0.3) axes1[0, 1].axhline(y=0, color="k", lw=0.5) axes1[1, 0].set_xlabel("Time t (s)") axes1[1, 0].set_ylabel("$u$ (torque)") axes1[1, 0].set_title("Control Input") axes1[1, 0].grid(True, alpha=0.3) axes1[1, 1].set_xlabel("Time t (s)") axes1[1, 1].set_ylabel("$V(x)$") axes1[1, 1].set_title("Value Function") axes1[1, 1].grid(True, alpha=0.3) fig1.suptitle("HPI Optimal Control — Closed-Loop Validation", fontsize=15, fontweight="bold") plt.tight_layout() path1 = os.path.join(OUTPUT_DIR, "trajectories.png") fig1.savefig(path1, dpi=150, bbox_inches="tight") print(f" Plot saved: {path1}") # ── Figure 2: Phase Portrait ── fig2, ax2 = plt.subplots(figsize=(8, 8)) for idx, (label, (t, X, U, V)) in enumerate(all_data.items()): c = colors[idx % len(colors)] ax2.plot(X[:, 0], X[:, 1], color=c, lw=1.5, label=label) ax2.scatter(X[0, 0], X[0, 1], color=c, s=80, marker="o", zorder=5) ax2.scatter(X[-1, 0], X[-1, 1], color=c, s=80, marker="x", zorder=5) ax2.set_xlabel("$x_1$ (angle)") ax2.set_ylabel("$x_2$ (angular velocity)") ax2.set_title("Phase Portrait") ax2.legend() ax2.grid(True, alpha=0.3) ax2.axhline(y=0, color="k", lw=0.5) ax2.axvline(x=0, color="k", lw=0.5) ax2.set_aspect("equal") path2 = os.path.join(OUTPUT_DIR, "phase_portrait.png") fig2.savefig(path2, dpi=150, bbox_inches="tight") print(f" Plot saved: {path2}") # ── Figure 3: Convergence ── fig3, axes3 = plt.subplots(1, 2, figsize=(14, 5)) L_vals = [h[0] for h in results["phase2_history"]] axes3[0].plot(range(len(L_vals)), L_vals, "o-", color="#1f77b4", markersize=6) axes3[0].set_xlabel("Iteration") axes3[0].set_ylabel("L") axes3[0].set_title("Phase 2: Homotopy Contraction") axes3[0].grid(True, alpha=0.3) loss_history = results["phase3_history"] if len(loss_history) > 1: loss_diffs = [abs(loss_history[i] - loss_history[i - 1]) for i in range(1, len(loss_history))] axes3[1].plot(range(len(loss_diffs)), loss_diffs, "o-", color="#d62728", markersize=6) axes3[1].set_yscale("log") axes3[1].set_xlabel("Iteration") axes3[1].set_ylabel("|loss change|") axes3[1].set_title("Phase 3: Loss Convergence (log)") axes3[1].grid(True, alpha=0.3) fig3.suptitle("HPI Convergence History", fontsize=14, fontweight="bold") plt.tight_layout() path3 = os.path.join(OUTPUT_DIR, "convergence.png") fig3.savefig(path3, dpi=150, bbox_inches="tight") print(f" Plot saved: {path3}") # ── Figure 4: Value function surface ── fig4, ax4 = plt.subplots(figsize=(8, 6)) x1_g = np.linspace(-1.2, 1.2, 80) x2_g = np.linspace(-2.5, 2.5, 80) X1, X2 = np.meshgrid(x1_g, x2_g) Xf = np.column_stack([X1.ravel(), X2.ravel()]) Vg = controller.trainer.compute_V(Xf).reshape(80, 80) vmin, vmax = max(Vg.min(), 1e-6), Vg.max() if vmax > vmin: levels = np.logspace(np.log10(vmin), np.log10(vmax), 20) cs = ax4.contourf(X1, X2, Vg, levels=levels, cmap="RdYlBu_r") ax4.contour(X1, X2, Vg, levels=levels[:8], colors="k", linewidths=0.3, alpha=0.5) fig4.colorbar(cs, ax=ax4, label="$V(x)$") for idx, (label, (t, X, U, V)) in enumerate(all_data.items()): ax4.plot(X[:, 0], X[:, 1], color=colors[idx % len(colors)], lw=1.5, label=label) ax4.set_xlabel("$x_1$ (angle)") ax4.set_ylabel("$x_2$ (ang. velocity)") ax4.set_title("Value Function with Phase Trajectories") ax4.legend(loc="upper left", fontsize=9) path4 = os.path.join(OUTPUT_DIR, "value_function.png") fig4.savefig(path4, dpi=150, bbox_inches="tight") print(f" Plot saved: {path4}") # ═══════════════════════════════════════════════════════════════════ # Main # ═══════════════════════════════════════════════════════════════════ def main(): # ── LQR baseline ── print("=" * 60) print(" LQR Baseline (linearized pendulum)") print("=" * 60) K = compute_lqr() print(f" LQR gain K = {K}") # ── HPI ── controller, results = run_hpi() # ── Validation ── print("\n" + "=" * 60) print(" Closed-Loop Validation") print("=" * 60) test_cases = [ ("IC1_small_angle", np.array([0.5, 0.0])), ("IC2_large_angle", np.array([0.9, 0.1])), ("IC3_neg_angle", np.array([-0.7, -0.2])), ("IC4_pos_vel", np.array([0.2, 1.5])), ("IC5_neg_vel", np.array([-0.3, -1.2])), ] all_data = {} summary = [] for label, x0 in test_cases: # HPI t_h, X_h, U_h, V_h = simulate_closed_loop(controller, x0) n_hpi = np.linalg.norm(X_h[-1]) # LQR t_l, X_l, U_l = simulate_lqr_closed_loop(controller.J, controller.mgl, K, x0) n_lqr = np.linalg.norm(X_l[-1]) all_data[f"HPI_{label}"] = (t_h, X_h, U_h, V_h) save_csv(t_h, X_h, U_h, V_h, f"hpi_{label}") print(f" {label:20s} HPI: ||x(T)||={n_hpi:.2e} | LQR: ||x(T)||={n_lqr:.2e}") summary.append({ "test_case": label, "x1_0": x0[0], "x2_0": x0[1], "hpi_norm_xT": n_hpi, "lqr_norm_xT": n_lqr, }) df_s = pd.DataFrame(summary) sp = os.path.join(OUTPUT_DIR, "summary.csv") df_s.to_csv(sp, index=False) print(f"\n Summary: {sp}") print(f"\n LQR K = {K}") # ── Plot ── print("\n Generating plots...") plot_results(all_data, controller, results) print("\n Done! Output files in ./output/") if __name__ == "__main__": main()