224 lines
9.4 KiB
Python
224 lines
9.4 KiB
Python
"""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
|