codex/optimizer-state-history.py

optimizer-state-history.py
"""Original deterministic momentum example; Python 3.14, NumPy/Matplotlib."""
from pathlib import Path
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

def trajectory(initial_buffer, steps=12):
    theta, buffer = 1.0, initial_buffer
    values = [theta]
    for _ in range(steps):
        buffer = .9 * buffer + theta
        theta -= .1 * buffer
        values.append(theta)
    return np.array(values)

retained, reset = trajectory(2.0), trajectory(0.0)
assert np.isclose(retained[1], .72) and np.isclose(reset[1], .9)
fig, ax = plt.subplots(figsize=(7.6, 3.5), layout="constrained", facecolor="white")
ax.plot(range(13), retained, "o-", markersize=4, lw=2, color="#26688c", label="Retained momentum buffer: b = 2")
ax.plot(range(13), reset, "s--", markersize=4, lw=2, color="#b45b26", label="Reset momentum buffer: b = 0")
ax.axhline(0, color="#777777", lw=.8)
ax.annotate("Same initial parameter: 1", xy=(0, 1), xytext=(3.5, .85), arrowprops={"arrowstyle":"->", "color":"#555555"}, fontsize=10)
ax.set(xlabel="Subsequent optimizer updates", ylabel="Parameter θ", title="Stored optimizer history changes the next step", xlim=(-.25, 12.25), ylim=(-.9, 1.15))
ax.set_xticks(range(0, 13, 2))
ax.grid(alpha=.2)
ax.legend(loc="lower left", fontsize=9)
fig.savefig(Path("optimizer-state-history.png"), dpi=140, facecolor="white")