first commit
This commit is contained in:
@@ -0,0 +1,26 @@
|
||||
"""HPI: Homotopy-Based Policy Iteration for Adaptive Optimal Control.
|
||||
|
||||
Implements the algorithm from:
|
||||
Chen et al., "Adaptive Optimal Control of Unknown Nonlinear Systems
|
||||
via Homotopy-Based Policy Iteration", IEEE TAC, 2024.
|
||||
|
||||
Uses neural network function approximators (CriticNN, ActorNN) trained
|
||||
via gradient-based Bellman residual minimization.
|
||||
"""
|
||||
|
||||
from .nn_models import (ActorNN, CriticNN, check_lyapunov_decrease,
|
||||
check_positive_definite, compute_q)
|
||||
from .data_collector import DataCollector
|
||||
from .nn_trainer import NNTrainer
|
||||
from .hpi_controller import HPIController
|
||||
|
||||
__all__ = [
|
||||
"CriticNN",
|
||||
"ActorNN",
|
||||
"compute_q",
|
||||
"check_positive_definite",
|
||||
"check_lyapunov_decrease",
|
||||
"DataCollector",
|
||||
"NNTrainer",
|
||||
"HPIController",
|
||||
]
|
||||
@@ -0,0 +1,131 @@
|
||||
"""Module A: ODE simulation and trajectory data collection.
|
||||
|
||||
Implements inverted pendulum dynamics and persistent excitation signal
|
||||
generation for data-driven adaptive optimal control.
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
from scipy.integrate import solve_ivp
|
||||
|
||||
|
||||
class DataCollector:
|
||||
"""Inverted pendulum simulator with data collection capabilities.
|
||||
|
||||
The pendulum dynamics are given by Eq. (27):
|
||||
dx₁/dt = x₂
|
||||
dx₂/dt = (mgl/J)·sin(x₁) + (1/J)·u
|
||||
"""
|
||||
|
||||
def __init__(self, J: float = 1.0, mgl: float = 1.0, n_freqs: int = 100):
|
||||
"""
|
||||
Args:
|
||||
J: Moment of inertia (default 1.0).
|
||||
mgl: Mass × gravity × length product (default 1.0).
|
||||
n_freqs: Number of frequency components in PE signal.
|
||||
"""
|
||||
self.J = float(J)
|
||||
self.mgl = float(mgl)
|
||||
self.n_freqs = n_freqs
|
||||
|
||||
# Pre-compute PE signal parameters
|
||||
self._rng = np.random.RandomState(42)
|
||||
self._omegas = self._rng.uniform(0.5, 50.0, size=n_freqs)
|
||||
self._a = self._rng.uniform(-1.0, 1.0, size=n_freqs)
|
||||
self._b = self._rng.uniform(-1.0, 1.0, size=n_freqs)
|
||||
# Scale amplitudes so that max |pe_signal| ≈ 0.1
|
||||
raw_max = np.sum(np.abs(self._a) + np.abs(self._b))
|
||||
self._scale = 0.1 / raw_max if raw_max > 0 else 1.0
|
||||
|
||||
def pe_signal(self, t):
|
||||
"""Persistent excitation signal.
|
||||
|
||||
u_e(t) = scale · Σ (a_k·sin(ω_k·t) + b_k·cos(ω_k·t))
|
||||
|
||||
Args:
|
||||
t: Time value (scalar).
|
||||
|
||||
Returns:
|
||||
float: PE signal value at time t.
|
||||
"""
|
||||
t = float(t)
|
||||
signals = self._a * np.sin(self._omegas * t) + self._b * np.cos(self._omegas * t)
|
||||
return self._scale * np.sum(signals)
|
||||
|
||||
def f(self, x):
|
||||
"""Drift dynamics f(x) — Eq. (27).
|
||||
|
||||
f(x) = [x₂, (mgl/J)·sin(x₁)]ᵀ
|
||||
|
||||
Args:
|
||||
x: State vector (2,) or (N, 2).
|
||||
|
||||
Returns:
|
||||
np.ndarray: Drift dynamics values.
|
||||
"""
|
||||
x = np.asarray(x, dtype=np.float64)
|
||||
if x.ndim == 1:
|
||||
return np.array([x[1], (self.mgl / self.J) * np.sin(x[0])])
|
||||
return np.column_stack([x[:, 1], (self.mgl / self.J) * np.sin(x[:, 0])])
|
||||
|
||||
def g(self, x):
|
||||
"""Input matrix g(x) — Eq. (27).
|
||||
|
||||
g(x) = [0, 1/J]ᵀ (constant)
|
||||
|
||||
Args:
|
||||
x: State vector (2,) or (N, 2) — ignored, kept for interface consistency.
|
||||
|
||||
Returns:
|
||||
np.ndarray: Input matrix [0, 1/J]ᵀ.
|
||||
"""
|
||||
x = np.asarray(x, dtype=np.float64)
|
||||
if x.ndim == 1:
|
||||
return np.array([0.0, 1.0 / self.J])
|
||||
N = x.shape[0]
|
||||
result = np.zeros((N, 2))
|
||||
result[:, 1] = 1.0 / self.J
|
||||
return result
|
||||
|
||||
def dynamics(self, t, x, u_func):
|
||||
"""ODE right-hand side: ẋ = f(x) + g(x)·u(t).
|
||||
|
||||
Args:
|
||||
t: Current time.
|
||||
x: State vector (2,).
|
||||
u_func: Callable u_func(t, x) → control input.
|
||||
|
||||
Returns:
|
||||
np.ndarray: State derivative (2,).
|
||||
"""
|
||||
u_val = u_func(t, x)
|
||||
return self.f(x) + self.g(x) * u_val
|
||||
|
||||
def collect_trajectory(self, x0, T, dt, u_func):
|
||||
"""Collect a trajectory by simulating the system.
|
||||
|
||||
Args:
|
||||
x0: Initial state (2,).
|
||||
T: Simulation duration.
|
||||
dt: Time step for output sampling.
|
||||
u_func: Control function u_func(t, x) → float.
|
||||
|
||||
Returns:
|
||||
tuple: (t, X, U) where
|
||||
t: Time points (M,).
|
||||
X: State trajectory (M, 2).
|
||||
U: Control inputs (M,).
|
||||
"""
|
||||
t_eval = np.arange(0, T, dt)
|
||||
t_span = (0.0, T)
|
||||
|
||||
def ode_rhs(t, x):
|
||||
return self.dynamics(t, x, u_func)
|
||||
|
||||
sol = solve_ivp(ode_rhs, t_span, x0, method="RK45",
|
||||
t_eval=t_eval, rtol=1e-8, atol=1e-10)
|
||||
|
||||
t = sol.t
|
||||
X = sol.y.T # (M, 2)
|
||||
U = np.array([u_func(ti, xi) for ti, xi in zip(t, X)])
|
||||
|
||||
return t, X, U
|
||||
+89
@@ -0,0 +1,89 @@
|
||||
"""Module E: Closed-loop verification demo for HPI algorithm — NN edition.
|
||||
|
||||
Runs the full three-phase HPI pipeline and verifies the optimal
|
||||
policy by simulating the closed-loop system from a perturbed initial state.
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
from scipy.integrate import solve_ivp
|
||||
|
||||
from .hpi_controller import HPIController
|
||||
|
||||
|
||||
def run_demo():
|
||||
"""Run full HPI demonstration and closed-loop verification.
|
||||
|
||||
Returns:
|
||||
dict: All results from the HPI pipeline and verification.
|
||||
"""
|
||||
print("=" * 60)
|
||||
print(" HPI: Homotopy-Based Policy Iteration (NN)")
|
||||
print(" Inverted Pendulum Adaptive Optimal Control")
|
||||
print("=" * 60)
|
||||
|
||||
# ── Initialize controller ──
|
||||
controller = HPIController(J=1.0, mgl=1.0, R=1.0, pe_amplitude=0.5,
|
||||
lr=1e-3, epochs=500, lambda_weight=1e-6)
|
||||
|
||||
# ── Run three-phase HPI ──
|
||||
print("\n[1/3] Phase 1: Searching for admissible L0 ...")
|
||||
results = controller.run(
|
||||
x0=np.array([0.5, 0.0]),
|
||||
T=5.0,
|
||||
dt=0.05,
|
||||
L_start=0.1,
|
||||
L_step=0.5,
|
||||
L_max=20.0,
|
||||
)
|
||||
|
||||
print(f" L0 = {results['phase1_L0']:.2f}")
|
||||
print(f" Phase 1 info: {results['phase1_info']}")
|
||||
print(f" Positive definite: {controller.check_positive_definite()}")
|
||||
|
||||
print(f"\n[2/3] Phase 2: Homotopy contraction ({len(results['phase2_history'])} iterations)")
|
||||
for i, (L, loss, pd) in enumerate(results["phase2_history"]):
|
||||
print(f" iter {i}: L={L:.4f}, loss={loss:.4e}, PD={pd}")
|
||||
|
||||
print(f"\n[3/3] Phase 3: Policy iteration (converged: {results['phase3_converged']})")
|
||||
print(f" Iterations: {len(results['phase3_history'])}")
|
||||
|
||||
# ── Closed-loop verification ──
|
||||
print("\n" + "=" * 60)
|
||||
print(" Closed-loop verification")
|
||||
print("=" * 60)
|
||||
|
||||
x0_test = np.array([0.9, 0.1])
|
||||
|
||||
def optimal_control(t, x):
|
||||
return controller.trainer.compute_u_hat(x)
|
||||
|
||||
def closed_loop_dynamics(t, x):
|
||||
return controller.collector.dynamics(t, x, optimal_control)
|
||||
|
||||
T_test = 5.0
|
||||
t_eval = np.linspace(0, T_test, 200)
|
||||
sol = solve_ivp(closed_loop_dynamics, (0, T_test), x0_test,
|
||||
method="RK45", t_eval=t_eval, rtol=1e-8, atol=1e-10)
|
||||
|
||||
X_cl = sol.y.T
|
||||
x_final = X_cl[-1]
|
||||
norm_final = np.linalg.norm(x_final)
|
||||
|
||||
print(f" Initial state: {x0_test}")
|
||||
print(f" Final state: {np.array2string(x_final, precision=8)}")
|
||||
print(f" ||x(T)||: {norm_final:.2e}")
|
||||
|
||||
if norm_final < 1e-3:
|
||||
print(" VERDICT: System converged to origin. [OK]")
|
||||
else:
|
||||
print(f" VERDICT: Residual norm {norm_final:.2e} > 1e-3 [FAIL]")
|
||||
|
||||
results["verification_x0"] = x0_test
|
||||
results["verification_x_final"] = x_final
|
||||
results["verification_norm"] = norm_final
|
||||
|
||||
return results
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
run_demo()
|
||||
@@ -0,0 +1,412 @@
|
||||
"""Module D: Three-phase HPI control flow (Algorithm 1) — NN edition.
|
||||
|
||||
Phase 1: Find admissible L0 with PD critic (NN-based).
|
||||
Phase 2: Homotopy contraction L_i -> 0 with safety constraint.
|
||||
Phase 3: Standard policy iteration on original system (L = 0).
|
||||
"""
|
||||
|
||||
import copy
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from .data_collector import DataCollector
|
||||
from .nn_models import check_lyapunov_decrease, check_positive_definite
|
||||
from .nn_trainer import NNTrainer
|
||||
|
||||
|
||||
class HPIController:
|
||||
"""Implements the three-phase Homotopy-Based Policy Iteration algorithm
|
||||
using neural network function approximators."""
|
||||
|
||||
def __init__(self, J=1.0, mgl=1.0, R=1.0, pe_amplitude=0.5,
|
||||
lr=1e-3, epochs=500, lambda_weight=1e-6):
|
||||
"""
|
||||
Args:
|
||||
J: Moment of inertia (default 1.0).
|
||||
mgl: Mass x gravity x length product (default 1.0).
|
||||
R: Control cost weight (default 1.0).
|
||||
pe_amplitude: PE signal amplitude scale (default 0.5).
|
||||
lr: Learning rate for Adam optimizer.
|
||||
epochs: Training epochs per HPI iteration.
|
||||
lambda_weight: L2 regularization coefficient.
|
||||
"""
|
||||
self.J = float(J)
|
||||
self.mgl = float(mgl)
|
||||
self.R = float(R)
|
||||
self.lr = float(lr)
|
||||
self.epochs = int(epochs)
|
||||
self.lambda_weight = float(lambda_weight)
|
||||
self.pe_amplitude = float(pe_amplitude)
|
||||
self.collector = DataCollector(J=J, mgl=mgl)
|
||||
self.trainer = NNTrainer(R=R, lr=lr, epochs=epochs,
|
||||
lambda_weight=lambda_weight)
|
||||
# Override PE amplitude
|
||||
self.collector._scale = pe_amplitude / np.sum(
|
||||
np.abs(self.collector._a) + np.abs(self.collector._b)
|
||||
)
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════
|
||||
# Utility
|
||||
# ═══════════════════════════════════════════════════════════════
|
||||
|
||||
def check_positive_definite(self, critic_nn=None, X_traj=None):
|
||||
"""Check if critic NN corresponds to a PD value function.
|
||||
|
||||
Checks both Hessian at origin and V(x) positivity on trajectory.
|
||||
"""
|
||||
if critic_nn is None:
|
||||
critic_nn = self.trainer.critic_nn
|
||||
return check_positive_definite(critic_nn, X_traj=X_traj)
|
||||
|
||||
def check_lyapunov_decrease(self, X, f_X):
|
||||
"""Check if current critic satisfies Lyapunov decrease on trajectory."""
|
||||
return check_lyapunov_decrease(self.trainer.critic_nn, X, f_X)
|
||||
|
||||
def _compute_f_X(self, X):
|
||||
"""Drift dynamics f(x) = [x2, sin(x1)] on trajectory."""
|
||||
X = np.asarray(X)
|
||||
f_X = np.zeros_like(X)
|
||||
f_X[:, 0] = X[:, 1]
|
||||
f_X[:, 1] = np.sin(X[:, 0])
|
||||
return f_X
|
||||
|
||||
def _hessian_info(self):
|
||||
"""Return (h11, h22, det) of the critic Hessian at origin for logging."""
|
||||
x0 = torch.zeros(1, 2, requires_grad=True)
|
||||
V0 = self.trainer.critic_nn(x0)
|
||||
grad_V = torch.autograd.grad(V0.sum(), x0, create_graph=True)[0]
|
||||
h11 = torch.autograd.grad(grad_V[0, 0], x0, retain_graph=True)[0][0, 0].item()
|
||||
h22 = torch.autograd.grad(grad_V[0, 1], x0, retain_graph=True)[0][0, 1].item()
|
||||
h12 = torch.autograd.grad(grad_V[0, 0], x0, retain_graph=True)[0][0, 1].item()
|
||||
det = h11 * h22 - h12 ** 2
|
||||
return h11, h22, det
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════
|
||||
# Phase 1: Find L0 — scan ALL L values, pick best loss
|
||||
# ═══════════════════════════════════════════════════════════════
|
||||
|
||||
def phase_one_find_L0(self, x0, T, dt, L_start=0.1, L_step=0.5, L_max=20.0,
|
||||
verbose=False):
|
||||
"""Find admissible L0 by scanning all L values and picking the best.
|
||||
|
||||
Unlike the polynomial version which returns at the FIRST PD L,
|
||||
this scans ALL L values and chooses the one with lowest Bellman
|
||||
residual loss AND positive-definite critic. This ensures Phase 2
|
||||
starts from the best possible critic.
|
||||
|
||||
Args:
|
||||
x0: Initial state (2,).
|
||||
T: Simulation duration.
|
||||
dt: Time step.
|
||||
L_start: Starting L value.
|
||||
L_step: Increment step.
|
||||
L_max: Maximum L value before failure.
|
||||
verbose: If True, print progress.
|
||||
|
||||
Returns:
|
||||
tuple: (L0, critic_sd, actor_sd, results_dict).
|
||||
|
||||
Raises:
|
||||
RuntimeError: If no admissible L is found below L_max.
|
||||
"""
|
||||
L_values = []
|
||||
L = L_start
|
||||
while L <= L_max:
|
||||
L_values.append(L)
|
||||
L += L_step
|
||||
|
||||
if not L_values:
|
||||
raise RuntimeError("No L values to try (L_start > L_max)")
|
||||
|
||||
# Collect ONE shared trajectory with PE-only control (no actor)
|
||||
# for training at all L values.
|
||||
def u_func(t, x):
|
||||
return self.collector.pe_signal(t)
|
||||
|
||||
t, X, U = self.collector.collect_trajectory(x0, T, dt, u_func)
|
||||
f_X = self._compute_f_X(X)
|
||||
|
||||
best_L = None
|
||||
best_loss = float('inf')
|
||||
best_critic_sd = None
|
||||
best_actor_sd = None
|
||||
best_info = None
|
||||
|
||||
for L in L_values:
|
||||
if verbose:
|
||||
print(f" Phase 1: trying L={L:.2f}...", flush=True)
|
||||
|
||||
# Fresh trainer for each L
|
||||
trainer = NNTrainer(R=self.R, lr=self.lr, epochs=self.epochs,
|
||||
lambda_weight=self.lambda_weight)
|
||||
|
||||
result = trainer.train_one_iteration(t, X, U, L, prev_actor=None, dt=dt)
|
||||
loss = result['loss_c']
|
||||
pd = check_positive_definite(trainer.critic_nn, X_traj=X)
|
||||
lyap = check_lyapunov_decrease(trainer.critic_nn, X, f_X)
|
||||
|
||||
if verbose:
|
||||
h11, h22, det = self._hessian_info_via_trainer(trainer)
|
||||
print(f" loss={loss:.4e} PD={pd} lyap={lyap:.3f} "
|
||||
f"H=[{h11:.4f},{h22:.4f}] det={det:.4e}", flush=True)
|
||||
|
||||
# Only consider PD critics; pick lowest loss
|
||||
if pd and loss < best_loss:
|
||||
best_L = L
|
||||
best_loss = loss
|
||||
best_critic_sd = {k: v.clone() for k, v in trainer.critic_nn.state_dict().items()}
|
||||
best_actor_sd = {k: v.clone() for k, v in trainer.actor_nn.state_dict().items()}
|
||||
best_info = {'loss': loss, 'lyap': lyap, 'mse': result['mse_c']}
|
||||
|
||||
if best_L is None:
|
||||
raise RuntimeError(
|
||||
f"Phase 1 failed: no admissible L found up to L_max={L_max}. "
|
||||
f"Try increasing L_max or adjusting training parameters."
|
||||
)
|
||||
|
||||
if verbose:
|
||||
print(f" Phase 1: selected L0={best_L:.2f} (loss={best_loss:.4e})",
|
||||
flush=True)
|
||||
|
||||
# Load best weights into main trainer
|
||||
self.trainer.critic_nn.load_state_dict(best_critic_sd)
|
||||
self.trainer.actor_nn.load_state_dict(best_actor_sd)
|
||||
|
||||
return best_L, best_critic_sd, best_actor_sd, best_info
|
||||
|
||||
def _hessian_info_via_trainer(self, trainer):
|
||||
"""Get Hessian info from a specific trainer instance."""
|
||||
x0 = torch.zeros(1, 2, requires_grad=True)
|
||||
V0 = trainer.critic_nn(x0)
|
||||
grad_V = torch.autograd.grad(V0.sum(), x0, create_graph=True)[0]
|
||||
h11 = torch.autograd.grad(grad_V[0, 0], x0, retain_graph=True)[0][0, 0].item()
|
||||
h22 = torch.autograd.grad(grad_V[0, 1], x0, retain_graph=True)[0][0, 1].item()
|
||||
h12 = torch.autograd.grad(grad_V[0, 0], x0, retain_graph=True)[0][0, 1].item()
|
||||
det = h11 * h22 - h12 ** 2
|
||||
return h11, h22, det
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════
|
||||
# Phase 2: Homotopy contraction
|
||||
# ═══════════════════════════════════════════════════════════════
|
||||
|
||||
def _compute_safe_alpha(self, X, prev_actor, L, gamma=0.1):
|
||||
"""Compute safe homotopy step alpha via inequality constraint."""
|
||||
x_tensor = torch.tensor(X, dtype=torch.float32, requires_grad=True)
|
||||
V = self.trainer.critic_nn(x_tensor)
|
||||
grad_V = torch.autograd.grad(V.sum(), x_tensor, create_graph=False)[0]
|
||||
dot_V = (grad_V * x_tensor).sum(dim=1).detach().numpy() # (M,)
|
||||
|
||||
Q = X[:, 0] ** 2 + X[:, 1] ** 2 # (M,)
|
||||
|
||||
if prev_actor is not None:
|
||||
with torch.no_grad():
|
||||
u_old = prev_actor(torch.tensor(X, dtype=torch.float32)).numpy()
|
||||
else:
|
||||
u_old = np.zeros(len(X))
|
||||
|
||||
with torch.no_grad():
|
||||
u_new = self.trainer.actor_nn(torch.tensor(X, dtype=torch.float32)).numpy()
|
||||
|
||||
gamma_c = 1.0 - gamma
|
||||
rhs = gamma_c * (Q + self.R * u_new ** 2) + self.R * (u_old - u_new) ** 2
|
||||
lhs = gamma_c * Q
|
||||
|
||||
eps = 1e-10
|
||||
lower_candidates = []
|
||||
upper_candidates = []
|
||||
|
||||
for k in range(len(dot_V)):
|
||||
d = dot_V[k]
|
||||
if d > eps:
|
||||
lower_candidates.append(lhs[k] / d)
|
||||
upper_candidates.append(rhs[k] / d)
|
||||
|
||||
if not lower_candidates or not upper_candidates:
|
||||
return None
|
||||
|
||||
alpha_lo = max(lower_candidates)
|
||||
alpha_hi = min(upper_candidates)
|
||||
|
||||
if alpha_lo <= alpha_hi and alpha_hi > 0:
|
||||
return min(alpha_hi, L)
|
||||
return None
|
||||
|
||||
def phase_two_contraction(self, x0, T, dt, L0, actor_sd_init, critic_sd_init,
|
||||
gamma=0.1, max_iter=50, verbose=False):
|
||||
"""Contract homotopy parameter L to zero via safety-constrained steps.
|
||||
|
||||
Returns:
|
||||
tuple: (actor_sd, critic_sd, history).
|
||||
"""
|
||||
L = L0
|
||||
self.trainer.critic_nn.load_state_dict(critic_sd_init)
|
||||
self.trainer.actor_nn.load_state_dict(actor_sd_init)
|
||||
self.trainer.reset_optimizer()
|
||||
|
||||
best_actor_sd = {k: v.clone() for k, v in actor_sd_init.items()}
|
||||
best_critic_sd = {k: v.clone() for k, v in critic_sd_init.items()}
|
||||
reject_streak = 0
|
||||
history = [(L, 0.0, True)] # (L, loss, PD)
|
||||
|
||||
for iteration in range(max_iter):
|
||||
if verbose:
|
||||
print(f" Phase 2 iter {iteration}: L={L:.4f}", flush=True)
|
||||
|
||||
prev_actor = self.trainer.copy_actor()
|
||||
|
||||
def u_func(t, x):
|
||||
pe = self.collector.pe_signal(t)
|
||||
u_policy = self.trainer.compute_u_hat(x)
|
||||
return pe + u_policy
|
||||
|
||||
t, X, U = self.collector.collect_trajectory(x0, T, dt, u_func)
|
||||
f_X = self._compute_f_X(X)
|
||||
|
||||
# Fresh optimizer per iteration
|
||||
self.trainer.reset_optimizer()
|
||||
result = self.trainer.train_one_iteration(t, X, U, L, prev_actor, dt)
|
||||
|
||||
pd_new = check_positive_definite(self.trainer.critic_nn, X_traj=X)
|
||||
alpha = self._compute_safe_alpha(X, prev_actor, L, gamma)
|
||||
|
||||
if alpha is None:
|
||||
alpha = L * 0.2
|
||||
|
||||
if pd_new:
|
||||
best_actor_sd = {k: v.clone() for k, v in self.trainer.actor_nn.state_dict().items()}
|
||||
best_critic_sd = {k: v.clone() for k, v in self.trainer.critic_nn.state_dict().items()}
|
||||
reject_streak = 0
|
||||
if verbose:
|
||||
lyap = check_lyapunov_decrease(self.trainer.critic_nn, X, f_X)
|
||||
print(f" alpha={alpha:.4f} loss={result['loss_c']:.4e} "
|
||||
f"lyap={lyap:.3f} PD", flush=True)
|
||||
else:
|
||||
self.trainer.actor_nn.load_state_dict(best_actor_sd)
|
||||
self.trainer.critic_nn.load_state_dict(best_critic_sd)
|
||||
reject_streak += 1
|
||||
if verbose:
|
||||
print(f" alpha={alpha:.4f} loss={result['loss_c']:.4e} "
|
||||
f"not PD (reject streak={reject_streak})", flush=True)
|
||||
|
||||
L = max(0.0, L - alpha)
|
||||
history.append((L, result['loss_c'], pd_new))
|
||||
|
||||
if L < 1e-14:
|
||||
if verbose:
|
||||
print(" L reached 0, Phase 2 complete", flush=True)
|
||||
break
|
||||
|
||||
if reject_streak > 15:
|
||||
if verbose:
|
||||
print(" >15 consecutive rejections, stopping Phase 2", flush=True)
|
||||
break
|
||||
|
||||
return best_actor_sd, best_critic_sd, history
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════
|
||||
# Phase 3: Standard policy iteration
|
||||
# ═══════════════════════════════════════════════════════════════
|
||||
|
||||
def phase_three_pi(self, x0, T, dt, actor_sd_init, critic_sd_init,
|
||||
epsilon=1e-6, max_iter=30, verbose=False):
|
||||
"""Standard policy iteration on the original system (L = 0).
|
||||
|
||||
Returns:
|
||||
tuple: (actor_sd, critic_sd, converged, history).
|
||||
"""
|
||||
self.trainer.actor_nn.load_state_dict(actor_sd_init)
|
||||
self.trainer.critic_nn.load_state_dict(critic_sd_init)
|
||||
self.trainer.reset_optimizer()
|
||||
|
||||
prev_loss = float('inf')
|
||||
converged = False
|
||||
history = []
|
||||
|
||||
for iteration in range(max_iter):
|
||||
if verbose:
|
||||
print(f" Phase 3 iter {iteration}", flush=True)
|
||||
|
||||
prev_actor = self.trainer.copy_actor()
|
||||
|
||||
def u_func(t, x):
|
||||
pe = self.collector.pe_signal(t)
|
||||
u_policy = self.trainer.compute_u_hat(x)
|
||||
u_policy = np.clip(u_policy, -5.0, 5.0)
|
||||
return pe + u_policy
|
||||
|
||||
t, X, U = self.collector.collect_trajectory(x0, T, dt, u_func)
|
||||
|
||||
self.trainer.reset_optimizer()
|
||||
result = self.trainer.train_one_iteration(t, X, U, 0.0, prev_actor, dt)
|
||||
loss = result['loss_c']
|
||||
history.append(loss)
|
||||
|
||||
pd = check_positive_definite(self.trainer.critic_nn, X_traj=X)
|
||||
|
||||
if verbose:
|
||||
f_X = self._compute_f_X(X)
|
||||
lyap = check_lyapunov_decrease(self.trainer.critic_nn, X, f_X)
|
||||
print(f" loss={loss:.4e} PD={pd} lyap={lyap:.3f}", flush=True)
|
||||
|
||||
if abs(prev_loss - loss) < epsilon:
|
||||
converged = True
|
||||
break
|
||||
|
||||
prev_loss = loss
|
||||
|
||||
if loss > 1e6 or not np.isfinite(loss):
|
||||
break
|
||||
|
||||
actor_sd = {k: v.clone() for k, v in self.trainer.actor_nn.state_dict().items()}
|
||||
critic_sd = {k: v.clone() for k, v in self.trainer.critic_nn.state_dict().items()}
|
||||
return actor_sd, critic_sd, converged, history
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════
|
||||
# Full pipeline
|
||||
# ═══════════════════════════════════════════════════════════════
|
||||
|
||||
def run(self, x0=None, T=10.0, dt=0.01,
|
||||
L_start=0.1, L_step=0.5, L_max=20.0,
|
||||
gamma=0.1, epsilon=1e-6, verbose=True):
|
||||
"""Run the full three-phase HPI algorithm.
|
||||
|
||||
Returns:
|
||||
dict: Results containing:
|
||||
- phase1_L0, phase1_critic_sd, phase1_actor_sd, phase1_info
|
||||
- phase2_history, phase2_actor_sd, phase2_critic_sd
|
||||
- phase3_actor_sd, phase3_critic_sd, phase3_converged, phase3_history
|
||||
"""
|
||||
if x0 is None:
|
||||
x0 = np.array([0.5, 0.0])
|
||||
|
||||
results = {}
|
||||
|
||||
# ── Phase 1: Find L0 ──
|
||||
L0, critic_sd1, actor_sd1, info1 = self.phase_one_find_L0(
|
||||
x0, T, dt, L_start, L_step, L_max, verbose=verbose)
|
||||
results["phase1_L0"] = L0
|
||||
results["phase1_critic_sd"] = critic_sd1
|
||||
results["phase1_actor_sd"] = actor_sd1
|
||||
results["phase1_info"] = info1
|
||||
|
||||
# ── Phase 2: Homotopy contraction ──
|
||||
actor_sd2, critic_sd2, hist2 = self.phase_two_contraction(
|
||||
x0, T, dt, L0, actor_sd1, critic_sd1, gamma=gamma, verbose=verbose)
|
||||
results["phase2_actor_sd"] = actor_sd2
|
||||
results["phase2_critic_sd"] = critic_sd2
|
||||
results["phase2_history"] = hist2
|
||||
|
||||
# ── Phase 3: Policy iteration ──
|
||||
actor_sd3, critic_sd3, converged, hist3 = self.phase_three_pi(
|
||||
x0, T, dt, actor_sd2, critic_sd2, epsilon, verbose=verbose)
|
||||
results["phase3_actor_sd"] = actor_sd3
|
||||
results["phase3_critic_sd"] = critic_sd3
|
||||
results["phase3_converged"] = converged
|
||||
results["phase3_history"] = hist3
|
||||
|
||||
# Load final weights
|
||||
self.trainer.actor_nn.load_state_dict(actor_sd3)
|
||||
self.trainer.critic_nn.load_state_dict(critic_sd3)
|
||||
|
||||
return results
|
||||
@@ -0,0 +1,147 @@
|
||||
"""Neural network models for HPI: CriticNN, ActorNN, and cost function.
|
||||
|
||||
Replaces the polynomial basis functions (symbolic_engine.py) with
|
||||
data-adaptive NN representations that avoid multicollinearity issues.
|
||||
"""
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class CriticNN(nn.Module):
|
||||
"""Value function approximator V(x) >= 0 with structural V(0)=0.
|
||||
|
||||
Architecture:
|
||||
Linear(2, 32, bias=False) -> Tanh
|
||||
Linear(32, 32, bias=False) -> Tanh
|
||||
Linear(32, 2, bias=False) # two raw outputs z1, z2
|
||||
V(x) = z1(x)^2 + z2(x)^2 # guaranteed V>=0 and V(0)=0
|
||||
|
||||
bias=False everywhere ensures V(0) = 0 structurally.
|
||||
Two output heads allow the Hessian at the origin to be rank-2 (strict PD possible).
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.fc1 = nn.Linear(2, 32, bias=False)
|
||||
self.fc2 = nn.Linear(32, 32, bias=False)
|
||||
self.fc3 = nn.Linear(32, 2, bias=False)
|
||||
|
||||
def forward(self, x):
|
||||
z = torch.tanh(self.fc1(x))
|
||||
z = torch.tanh(self.fc2(z))
|
||||
z = self.fc3(z) # (..., 2): [z1, z2]
|
||||
return (z[:, 0:1] ** 2 + z[:, 1:2] ** 2).squeeze(-1) # (...,) scalar
|
||||
|
||||
def get_raw_z(self, x):
|
||||
"""Return the two raw outputs z1, z2 before squaring."""
|
||||
z = torch.tanh(self.fc1(x))
|
||||
z = torch.tanh(self.fc2(z))
|
||||
return self.fc3(z) # (..., 2)
|
||||
|
||||
|
||||
class ActorNN(nn.Module):
|
||||
"""Policy network u(x): R^2 -> R.
|
||||
|
||||
Architecture:
|
||||
Linear(2, 32, bias=True) -> Tanh
|
||||
Linear(32, 32, bias=True) -> Tanh
|
||||
Linear(32, 1, bias=True) # signed control output
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.fc1 = nn.Linear(2, 32, bias=True)
|
||||
self.fc2 = nn.Linear(32, 32, bias=True)
|
||||
self.fc3 = nn.Linear(32, 1, bias=True)
|
||||
|
||||
def forward(self, x):
|
||||
h = torch.tanh(self.fc1(x))
|
||||
h = torch.tanh(self.fc2(h))
|
||||
return self.fc3(h).squeeze(-1) # (...,) scalar
|
||||
|
||||
|
||||
def compute_q(x):
|
||||
"""Cost function Q(x) = x1^2 + x2^2.
|
||||
|
||||
Args:
|
||||
x: torch.Tensor of shape (..., 2).
|
||||
|
||||
Returns:
|
||||
torch.Tensor of shape (...,).
|
||||
"""
|
||||
return x[..., 0] ** 2 + x[..., 1] ** 2
|
||||
|
||||
|
||||
def check_positive_definite(critic_nn, X_traj=None):
|
||||
"""Check if the critic NN corresponds to a positive-definite value function.
|
||||
|
||||
Two checks:
|
||||
1. Hessian at origin is strictly positive definite (necessary condition).
|
||||
2. V(x) > 0 on trajectory data (sanity check, structural guarantee makes
|
||||
this always true, but verifies the NN hasn't degenerated).
|
||||
|
||||
Args:
|
||||
critic_nn: CriticNN instance.
|
||||
X_traj: Optional trajectory data (M, 2) numpy array for additional check.
|
||||
|
||||
Returns:
|
||||
bool: True if all checks pass.
|
||||
"""
|
||||
# 1. Hessian at origin
|
||||
x0 = torch.zeros(1, 2, requires_grad=True)
|
||||
V0 = critic_nn(x0)
|
||||
grad_V = torch.autograd.grad(V0.sum(), x0, create_graph=True)[0] # (1, 2)
|
||||
|
||||
h11 = torch.autograd.grad(grad_V[0, 0], x0, retain_graph=True)[0][0, 0].item()
|
||||
h22 = torch.autograd.grad(grad_V[0, 1], x0, retain_graph=True)[0][0, 1].item()
|
||||
h12 = torch.autograd.grad(grad_V[0, 0], x0, retain_graph=True)[0][0, 1].item()
|
||||
|
||||
hessian_pd = bool(h11 > 1e-6 and h22 > 1e-6 and h11 * h22 - h12 ** 2 > 1e-12)
|
||||
|
||||
if not hessian_pd:
|
||||
return False
|
||||
|
||||
# 2. Optional: verify V(x) is well-behaved on trajectory data
|
||||
if X_traj is not None:
|
||||
x_t = torch.tensor(X_traj, dtype=torch.float32)
|
||||
with torch.no_grad():
|
||||
V_traj = critic_nn(x_t)
|
||||
if torch.any(torch.isnan(V_traj)) or torch.any(torch.isinf(V_traj)):
|
||||
return False
|
||||
# V should be finite and non-negative everywhere on trajectory
|
||||
if torch.any(V_traj < -1e-4):
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
|
||||
def check_lyapunov_decrease(critic_nn, X, f_X):
|
||||
"""Check if V satisfies the Lyapunov decrease condition on trajectory data.
|
||||
|
||||
Verifies: dV/dx · f(x) < 0 for states away from origin.
|
||||
This is a necessary condition for V to be a valid Lyapunov function.
|
||||
|
||||
Args:
|
||||
critic_nn: CriticNN instance.
|
||||
X: State trajectory (M, 2) numpy array.
|
||||
f_X: Drift dynamics f(X) evaluated on trajectory, (M, 2) numpy array.
|
||||
|
||||
Returns:
|
||||
float: Fraction of points (away from origin) where V_dot < 0.
|
||||
"""
|
||||
x_t = torch.tensor(X, dtype=torch.float32, requires_grad=True)
|
||||
f_t = torch.tensor(f_X, dtype=torch.float32)
|
||||
V = critic_nn(x_t)
|
||||
grad_V = torch.autograd.grad(V.sum(), x_t, create_graph=False)[0]
|
||||
V_dot = (grad_V * f_t).sum(dim=1) # (M,)
|
||||
|
||||
# Only check points away from origin (||x|| > 0.05)
|
||||
norms = torch.norm(x_t, dim=1)
|
||||
mask = norms > 0.05
|
||||
if mask.sum() == 0:
|
||||
return 1.0
|
||||
|
||||
V_dot_filtered = V_dot[mask]
|
||||
fraction_negative = (V_dot_filtered < 0).float().mean().item()
|
||||
return fraction_negative
|
||||
@@ -0,0 +1,223 @@
|
||||
"""NN-based HPI trainer: gradient-based Bellman residual minimization.
|
||||
|
||||
Replaces the linear-system solver (solver.py) with Adam optimization of
|
||||
neural network critic and actor parameters.
|
||||
|
||||
Loss functions:
|
||||
Critic: Bellman residual (policy evaluation)
|
||||
Actor: Supervised regression to optimal control law
|
||||
u*(x) = -1/(2R) * g^T(x) * grad_V(x)
|
||||
For pendulum: u*(x) = -0.5 * dV/dx2 (with J=1, R=1)
|
||||
"""
|
||||
|
||||
import copy
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from .nn_models import ActorNN, CriticNN
|
||||
|
||||
|
||||
class NNTrainer:
|
||||
"""Trains CriticNN and ActorNN via decoupled loss functions.
|
||||
|
||||
Critic: Bellman residual minimization (evaluated with current policy,
|
||||
but actor gradients blocked).
|
||||
Actor: Supervised regression to the greedy optimal policy
|
||||
u* = -1/(2R) * g^T * grad_V.
|
||||
"""
|
||||
|
||||
def __init__(self, R=1.0, J=1.0, lr=1e-3, epochs=300, lambda_weight=1e-6):
|
||||
self.R = float(R)
|
||||
self.J = float(J)
|
||||
self.lr = float(lr)
|
||||
self.epochs = int(epochs)
|
||||
self.lambda_weight = float(lambda_weight)
|
||||
|
||||
self.critic_nn = CriticNN()
|
||||
self.actor_nn = ActorNN()
|
||||
self._create_optimizers()
|
||||
|
||||
def _create_optimizers(self):
|
||||
self.optimizer_c = torch.optim.Adam(self.critic_nn.parameters(), lr=self.lr)
|
||||
self.optimizer_a = torch.optim.Adam(self.actor_nn.parameters(), lr=self.lr)
|
||||
|
||||
def reset_optimizer(self):
|
||||
"""Reset Adam optimizer states — call between HPI iterations."""
|
||||
self._create_optimizers()
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════
|
||||
# Forward evaluation (numpy -> torch -> numpy)
|
||||
# ═══════════════════════════════════════════════════════════════
|
||||
|
||||
def _to_tensor(self, x):
|
||||
x = np.asarray(x, dtype=np.float32)
|
||||
if x.ndim == 1:
|
||||
x = x.reshape(1, -1)
|
||||
return torch.from_numpy(x)
|
||||
|
||||
def compute_u_hat(self, x_np):
|
||||
"""Compute estimated optimal control u_hat = actor_nn(x)."""
|
||||
was_1d = np.asarray(x_np).ndim == 1
|
||||
x_tensor = self._to_tensor(x_np)
|
||||
with torch.no_grad():
|
||||
u = self.actor_nn(x_tensor).numpy()
|
||||
if was_1d:
|
||||
return float(u[0])
|
||||
return u
|
||||
|
||||
def compute_v_hat(self, u_np, x_np):
|
||||
"""Compute control residual: v_hat = u - u_hat."""
|
||||
u_hat = self.compute_u_hat(x_np)
|
||||
return np.asarray(u_np, dtype=np.float64) - np.asarray(u_hat, dtype=np.float64)
|
||||
|
||||
def compute_V(self, x_np):
|
||||
"""Compute value function V(x) = critic_nn(x)."""
|
||||
was_1d = np.asarray(x_np).ndim == 1
|
||||
x_tensor = self._to_tensor(x_np)
|
||||
with torch.no_grad():
|
||||
v = self.critic_nn(x_tensor).numpy()
|
||||
if was_1d:
|
||||
return float(v[0])
|
||||
return v
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════
|
||||
# Actor management
|
||||
# ═══════════════════════════════════════════════════════════════
|
||||
|
||||
def copy_actor(self):
|
||||
"""Return a deep copy of the current actor network."""
|
||||
return copy.deepcopy(self.actor_nn)
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════
|
||||
# Integration
|
||||
# ═══════════════════════════════════════════════════════════════
|
||||
|
||||
@staticmethod
|
||||
def integrate_interval(f_vals, dt):
|
||||
"""Trapezoidal integration over intervals."""
|
||||
f_vals = np.asarray(f_vals)
|
||||
return 0.5 * dt * (f_vals[:-1] + f_vals[1:])
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════
|
||||
# Training
|
||||
# ═══════════════════════════════════════════════════════════════
|
||||
|
||||
def train_one_iteration(self, t, X, U, L_i, prev_actor, dt):
|
||||
"""Train critic and actor NNs for one HPI iteration.
|
||||
|
||||
Decoupled training:
|
||||
1. Critic loss: Bellman residual r_k = DeltaV_k - I_grad_Lx_k + I_qu_k + I_au_k
|
||||
Actor output is detached so gradients flow only to the critic.
|
||||
2. Actor loss: supervised regression to u* = -1/(2R) * g^T * grad_V
|
||||
Critic gradient is detached so gradients flow only to the actor.
|
||||
|
||||
Args:
|
||||
t: Time points (M,).
|
||||
X: State trajectory (M, 2).
|
||||
U: Control inputs (M,) = PE + old_policy.
|
||||
L_i: Current homotopy parameter.
|
||||
prev_actor: Frozen ActorNN from previous iteration (or None for u_old=0).
|
||||
dt: Time step.
|
||||
|
||||
Returns:
|
||||
dict: training metrics.
|
||||
"""
|
||||
if dt is None:
|
||||
dt = t[1] - t[0]
|
||||
|
||||
x_tensor = torch.tensor(X, dtype=torch.float32)
|
||||
u_tensor = torch.tensor(U, dtype=torch.float32)
|
||||
|
||||
initial_loss_c = None
|
||||
initial_loss_a = None
|
||||
|
||||
for epoch in range(self.epochs):
|
||||
# ── Common forward pass for critic ──
|
||||
x_grad = x_tensor.clone().requires_grad_(True)
|
||||
V = self.critic_nn(x_grad) # (M,)
|
||||
grad_V = torch.autograd.grad(V.sum(), x_grad, create_graph=True)[0] # (M, 2)
|
||||
dV_dx_x = (grad_V * x_grad).sum(dim=1) # (M,)
|
||||
|
||||
# Target for actor: u* = -1/(2R) * g^T * grad_V
|
||||
# Pendulum: g(x) = [0, 1/J]^T, so u* = -1/(2*R*J) * dV/dx2
|
||||
u_target = (-0.5 / (self.R * self.J) * grad_V[:, 1]).detach()
|
||||
|
||||
# Current actor output and old actor output
|
||||
u_hat_new = self.actor_nn(x_tensor) # (M,) — gets actor gradients
|
||||
if prev_actor is not None:
|
||||
with torch.no_grad():
|
||||
u_hat_old = prev_actor(x_tensor) # (M,)
|
||||
else:
|
||||
u_hat_old = torch.zeros_like(u_hat_new)
|
||||
|
||||
# Q values
|
||||
Q = x_tensor[:, 0] ** 2 + x_tensor[:, 1] ** 2 # (M,)
|
||||
|
||||
# ── 1. Critic loss: Bellman residual ──
|
||||
# ΔV_k = V(x_{k+1}) - V(x_k)
|
||||
delta_V = V[1:] - V[:-1] # (M-1,)
|
||||
|
||||
# I_grad_Lx = (dt/2) * L * [gradV_k·x_k + gradV_{k+1}·x_{k+1}]
|
||||
I_grad_Lx = 0.5 * dt * L_i * (dV_dx_x[:-1] + dV_dx_x[1:]) # (M-1,)
|
||||
|
||||
# I_qu = (dt/2) * [(Q_k+R*u_old_k^2) + (Q_{k+1}+R*u_old_{k+1}^2)]
|
||||
qu = Q + self.R * u_hat_old ** 2 # (M,)
|
||||
I_qu = 0.5 * dt * (qu[:-1] + qu[1:]) # (M-1,)
|
||||
|
||||
# I_au = (dt/2) * [2R*u_new_k*(u_k-u_old_k) + 2R*u_new_{k+1}*(u_{k+1}-u_old_{k+1})]
|
||||
# Use detached u_hat_new so critic doesn't get actor gradients
|
||||
au = 2.0 * self.R * u_hat_new.detach() * (u_tensor - u_hat_old) # (M,)
|
||||
I_au = 0.5 * dt * (au[:-1] + au[1:]) # (M-1,)
|
||||
|
||||
r_c = delta_V - I_grad_Lx + I_qu + I_au # (M-1,)
|
||||
mse_c = torch.mean(r_c ** 2)
|
||||
l2_c = sum(p.pow(2.0).sum() for p in self.critic_nn.parameters())
|
||||
loss_c = mse_c + self.lambda_weight * l2_c
|
||||
|
||||
self.optimizer_c.zero_grad()
|
||||
loss_c.backward()
|
||||
torch.nn.utils.clip_grad_norm_(self.critic_nn.parameters(), max_norm=10.0)
|
||||
self.optimizer_c.step()
|
||||
|
||||
# ── 2. Actor loss: supervised regression to optimal policy ──
|
||||
mse_a = torch.mean((u_hat_new - u_target) ** 2)
|
||||
l2_a = sum(p.pow(2.0).sum() for p in self.actor_nn.parameters())
|
||||
loss_a = mse_a + self.lambda_weight * l2_a
|
||||
|
||||
self.optimizer_a.zero_grad()
|
||||
loss_a.backward()
|
||||
torch.nn.utils.clip_grad_norm_(self.actor_nn.parameters(), max_norm=10.0)
|
||||
self.optimizer_a.step()
|
||||
|
||||
if initial_loss_c is None:
|
||||
initial_loss_c = float(loss_c.item())
|
||||
initial_loss_a = float(loss_a.item())
|
||||
|
||||
# ── Lyapunov check on final critic ──
|
||||
x_grad_final = x_tensor.clone().requires_grad_(True)
|
||||
V_final = self.critic_nn(x_grad_final)
|
||||
grad_V_final = torch.autograd.grad(V_final.sum(), x_grad_final)[0]
|
||||
with torch.no_grad():
|
||||
f_t = torch.tensor(self._compute_f_X(X), dtype=torch.float32)
|
||||
V_dot = (grad_V_final * f_t).sum(dim=1)
|
||||
norms = torch.norm(x_tensor, dim=1)
|
||||
mask = norms > 0.05
|
||||
lyap_frac = (V_dot[mask] < 0).float().mean().item() if mask.sum() > 0 else 1.0
|
||||
|
||||
return {
|
||||
"loss_c": float(loss_c.item()),
|
||||
"mse_c": float(mse_c.item()),
|
||||
"loss_a": float(loss_a.item()),
|
||||
"mse_a": float(mse_a.item()),
|
||||
"initial_loss_c": initial_loss_c or float(loss_c.item()),
|
||||
"lyap_frac": lyap_frac,
|
||||
}
|
||||
|
||||
def _compute_f_X(self, X):
|
||||
"""Compute drift dynamics f(x) = [x2, sin(x1)] on trajectory."""
|
||||
X = np.asarray(X)
|
||||
f_X = np.zeros_like(X)
|
||||
f_X[:, 0] = X[:, 1]
|
||||
f_X[:, 1] = np.sin(X[:, 0])
|
||||
return f_X
|
||||
Reference in New Issue
Block a user