"""Generate cumulative regret plot for the Volterra memory example.""" import functools import numpy as np import flax.linen as nn import jax.numpy as jnp import jax.random as jr from jax import config, jit, value_and_grad from jax.lax import scan from jax.scipy.linalg import expm import matplotlib.pyplot as plt import optax from tqdm import tqdm from mpm import get_system_params, legsval, whitesignal, wray_and_green_output def log_downsample(x, y, num=2000): idx = jnp.unique(jnp.logspace(0, jnp.log10(len(x)-1), num).astype(int)) return x[idx], y[idx] def initialize_predictor(n: int, *, seed: int = 0): """Initialize parameters and optimizer for the linear/quadratic predictors.""" key = jr.PRNGKey(seed) key1, key2 = jr.split(key) params = { "w1": 1e-2 * jr.normal(key1, (n,)), "w2": 0.0 * jr.normal(key2, (n, n)), "b": jnp.zeros(1), } optimizer = optax.sgd(learning_rate=3e-2) opt_state = optimizer.init(params) return params, optimizer, opt_state def initialize_mlp_predictor(n: int, *, hidden_dim: int = 128, seed: int = 0): """Initialize parameters and optimizer for the MLP predictor.""" key = jr.PRNGKey(seed) model = MlpReadout(hidden_dim=hidden_dim) params = model.init(key, jnp.zeros((n,))) optimizer = optax.sgd(learning_rate=1e-2) opt_state = optimizer.init(params) return model, params, optimizer, opt_state class MlpReadout(nn.Module): hidden_dim: int = 128 @nn.compact def __call__(self, state): hidden = nn.Dense(self.hidden_dim)(state) hidden = jnp.tanh(hidden) pred = nn.Dense(1)(hidden) return jnp.squeeze(pred, axis=-1) def loss_func_linear(params, state, out_signal): w1, b = params["w1"], params["b"] pred = b[0] + jnp.inner(state, w1) return jnp.mean((pred - out_signal) ** 2), pred def loss_func_quadratic(params, state, out_signal): w1, w2, b = params["w1"], params["w2"], params["b"] pred = b[0] + jnp.inner(state, w1) pred = pred + jnp.sum(w2 * jnp.outer(state, state)) return jnp.mean((pred - out_signal) ** 2), pred def loss_func_mlp(params, state, out_signal, *, model): pred = model.apply(params, state) return jnp.mean((pred - out_signal) ** 2), pred def run_regret_experiment(): dt = 1e-2 # Build the HIPPO system. n = 64 measure = "legs" params, _, _ = get_system_params(measure, n) A, b = params timescale = 3 * dt / 0.08 A, b = A / timescale, b / timescale A_d = expm(dt * A) b_d = jnp.linalg.solve(A, A_d @ b - b) A, b = jnp.asarray(A_d), jnp.asarray(b_d) loss_funcs = [loss_func_linear, loss_func_quadratic, None] labels = ["Linear", "Quadratic", "MLP"] colors = ["peru", "mediumseagreen", "steelblue"] fig = plt.figure(figsize=(7, 2.5)) ax_regret = fig.add_subplot(1, 3, 1) plt.sca(ax_regret) quadratic_params = None num_trials = 5 all_cumulative = {label: [] for label in labels} pbar = tqdm(range(num_trials)) for trial_idx in pbar: np.random.seed(trial_idx) input_data = jnp.asarray(whitesignal(1e7 * dt, dt, freq=10)) output_data = jnp.asarray(wray_and_green_output(np.asarray(input_data))) signals = jnp.stack([input_data, output_data], axis=1) for iter_idx, (loss_func, label, color) in enumerate(zip(loss_funcs, labels, colors)): seed = trial_idx * len(labels) + iter_idx if label == "MLP": model, params, optimizer, opt_state = initialize_mlp_predictor(n, seed=seed) loss_func_i = functools.partial(loss_func_mlp, model=model) else: params, optimizer, opt_state = initialize_predictor(n, seed=seed) loss_func_i = loss_func val_grad_loss = jit(value_and_grad(loss_func_i, has_aux=True)) def step(carry, signals): in_x, out_x = signals state, params, opt_state = carry # Evolve the HIPPO state. state = A @ state + in_x * b # Calculate loss and update predictor. (loss_val, pred), grads = val_grad_loss(params, state, out_x) updates, opt_state = optimizer.update(grads, opt_state) params = optax.apply_updates(params, updates) return (state, params, opt_state), (loss_val, pred) initial = (jnp.zeros(n), params, opt_state) (_, params, _), (loss_vals, _) = scan(step, initial, signals) all_cumulative[label].append(jnp.cumsum(loss_vals)) if label == "Quadratic" and trial_idx == 0: quadratic_params = params steps = jnp.arange(len(all_cumulative[labels[0]][0])) for label, color in zip(labels, colors): curves = jnp.stack(all_cumulative[label]) mean = jnp.mean(curves, axis=0) sem = jnp.std(curves, axis=0) / jnp.sqrt(num_trials) x_ds, mean_ds = log_downsample(steps, mean) _, sem_ds = log_downsample(steps, sem) ax_regret.loglog(x_ds, mean_ds, label=label, c=color) ax_regret.fill_between(x_ds, mean_ds - sem_ds, mean_ds + sem_ds, color=color, alpha=0.3) plt.xlabel("Step") plt.xlim(10, None) plt.ylim(1e-1, None) plt.ylabel("Cumulative Error") plt.legend() # Visualize the true Volterra kernel. ax = fig.add_subplot(1, 3, 2) plt.sca(ax) tau_vals = dt * jnp.arange(50) a, m, k = 2.0, 0.3, 0.08 mu = lambda t: a / m * jnp.exp(-k * t) * jnp.sin(m * t) true_filter = mu(jnp.arange(50)) true_kernel = 4e-3 * jnp.outer(true_filter, true_filter) vmax = 0.06 plt.imshow( true_kernel, cmap="coolwarm", extent=[0, 49, 0, 49], origin="lower", vmax=vmax, vmin=-vmax, ) ax.set_title("True Kernel") ax.set_xlabel(r"$\tau_1$") ax.set_ylabel(r"$\tau_2$") # Visualize the inferred Volterra kernel from the quadratic predictor. ax = fig.add_subplot(1, 3, 3) plt.sca(ax) if quadratic_params is not None: w2 = quadratic_params["w2"] p_tau = dt * 1.0 / timescale * jnp.exp(-1.0 / timescale * tau_vals) legsvals = jnp.zeros((len(tau_vals), n)) for i in range(n): c = np.zeros(n) c[i] = 1.0 legsvals = legsvals.at[:, i].set(legsval(np.asarray(tau_vals / timescale), c)) legsvals = legsvals[None, :, None] * legsvals[:, None, :, None] legsvals = ( legsvals * w2[None, None] * p_tau.reshape(1, -1, 1, 1) * p_tau.reshape(-1, 1, 1, 1) ) kernel = jnp.sum(legsvals, axis=(-1, -2)) else: kernel = jnp.zeros((50, 50)) plt.imshow( kernel, cmap="coolwarm", extent=[0, 49, 0, 49], origin="lower", vmax=vmax, vmin=-vmax, ) ax.set_title("Inferred Kernel") ax.set_xlabel(r"$\tau_1$") ax.set_ylabel(r"$\tau_2$") plt.tight_layout() plt.savefig("volterra_regret.png") plt.savefig("volterra_regret.pdf") plt.close("all") if __name__ == "__main__": run_regret_experiment()