auto_math_lyapunov_stability_5

Snippet from Math-Lyapunov-Stability.md

import numpy as np

from compitum.control import LyapunovController

# The real LyapunovController has no accept/reject gate and no
# initial_radius/ema_alpha constructor params -- it always executes an
# update and returns (eta_cap, status). eta_cap = kappa / (grad_norm + 1e-6)
# is a real step-size cap; this adapts the Rosenbrock demo to actually use
# it (scaling the learning rate) instead of a fictional accept/reject gate.


def rosenbrock(x, y):
    """Rosenbrock function (non-convex test case)."""
    return (1 - x) ** 2 + 100 * (y - x**2) ** 2


def noisy_gradient(x, y, rng, noise_std=0.1):
    """Gradient with additive noise."""
    grad = np.array([
        -2 * (1 - x) - 400 * x * (y - x**2),
        200 * (y - x**2),
    ])
    noise = rng.normal(0, noise_std, size=2)
    return grad + noise


# Initialize
rng = np.random.default_rng(42)
x_traj = [np.array([-0.5, 2.0])]  # Start away from minimum (1, 1)
controller = LyapunovController(kappa=0.1, r0=0.5, integral_gain=0.005)
base_lr = 0.001

for step in range(200):
    x = x_traj[-1]
    E_t = rosenbrock(x[0], x[1])

    grad = noisy_gradient(x[0], x[1], rng)
    grad_norm = float(np.linalg.norm(grad))

    x_proposal = x - base_lr * grad
    E_tp1 = rosenbrock(x_proposal[0], x_proposal[1])
    # Rosenbrock's energy scale is large (thousands, away from the minimum);
    # the real controller's drift_integral accumulates d_star every step and
    # squares it in lyapunov_function() -- feeding it raw energy deltas
    # overflows within ~200 steps. d_star is meant to be a distance-like
    # quantity (see energy.py), not a raw loss delta, so clip it to a
    # reasonable range for this controller to stay well-posed.
    d_star = min(max(E_tp1 - E_t, 0.0), 5.0)

    _, status = controller.update(d_star, grad_norm)
    # Scale this step's effective learning rate by the controller's current
    # trust_radius, normalized against its initial r0 -- shrinks the step
    # when drift/uncertainty grows, expands it when things are stable. Capped
    # at 2x base_lr: trust_radius can grow past r0, and on a surface as steep
    # as Rosenbrock's, an uncapped multiplier on an already-large gradient
    # diverges within a handful of steps.
    effective_lr = base_lr * min(status["trust_radius"] / 0.5, 2.0)
    x_traj.append(x - effective_lr * grad)

print(f"Final position: ({x_traj[-1][0]:.3f}, {x_traj[-1][1]:.3f})")
print(f"Final energy: {rosenbrock(x_traj[-1][0], x_traj[-1][1]):.3f}")
print(f"Final trust_radius: {controller.trust_radius:.4f}")
print(f"Final drift_ema: {controller.drift_ema:.4f}")