413 lines
17 KiB
Python
413 lines
17 KiB
Python
"""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
|