330 lines
12 KiB
Python
330 lines
12 KiB
Python
"""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()
|