"""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