""" Generate a 2-panel figure illustrating REINFORCE variance scaling with trajectory length T: Var(g) = O(T^3). Setup: Multi-step environment with softmax linear policy. - State features: s_t ~ Uniform([-1, 1]^n), with fixed n = 4. - Policy: softmax linear pi_theta(a | s) = exp(theta_a^T s) / sum_a' exp(theta_{a'}^T s) - Number of actions: |A| = 4 (fixed) - theta initialized near zero - Reward per step: r(s_t, a_t) = 1 + 0.1 * (theta_{a_t}^T s_t) (mean reward ~ 1) - Score: X_t = (e_{a_t} - pi_t) ⊗ s_t - Estimator: g(tau) = R(tau) * S(tau) where S = sum_t X_t """ import numpy as np import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt from scipy.special import softmax # Color palette (Dark theme / Nord palette) BG_COLOR_FIG = "#1E222A" # Outer figure background BG_COLOR_AX = "#252931" # Subplot panel background COLOR_TEXT = "#ECEFF4" # Primary text color COLOR_MUTED = "#D8DEE9" # Muted text color GRID_COLOR = "#3B4252" # Grid line color SPINE_COLOR = "#4C566A" # Border / spine color COLOR_T = "#88C0D0" # Nord cyan/blue COLOR_RATIO = "#A3BE8C" # Nord green COLOR_REF = "#E5E9F0" # Dotted reference line def simulate_reinforce_gradients(n_samples, T, n_feat, n_actions, rng): """ Simulates trajectories and returns the REINFORCE gradient estimates. The policy theta is kept identical using a fixed seed generator. """ theta_rng = np.random.default_rng(0) theta = theta_rng.normal(0, 0.01, (n_actions, n_feat)) d = n_actions * n_feat total_scores = np.zeros((n_samples, d)) total_rewards = np.zeros(n_samples) for t in range(T): # Sample states: (n_samples, n_feat) states = rng.uniform(-1, 1, (n_samples, n_feat)) # Compute logits and action probabilities: (n_samples, n_actions) logits = states @ theta.T probs = softmax(logits, axis=1) # Sample actions using cumulative probability thresholding cum_probs = np.cumsum(probs, axis=1) u = rng.uniform(0, 1, n_samples) actions = np.sum(cum_probs < u[:, None], axis=1) actions = np.clip(actions, 0, n_actions - 1) # Compute scores: (e_{a_t} - pi_t) ⊗ s_t e_a = np.zeros((n_samples, n_actions)) e_a[np.arange(n_samples), actions] = 1.0 diff = e_a - probs # Kronecker product: each row of diff ⊗ each row of states score_vecs = diff[:, :, None] * states[:, None, :] score_vecs = score_vecs.reshape(n_samples, d) total_scores += score_vecs # Compute step rewards: 1.0 + 0.1 * (theta_a^T s_t) theta_a = theta[actions] rewards = 1.0 + 0.1 * np.sum(theta_a * states, axis=1) total_rewards += rewards # REINFORCE estimator: g = R(tau) * S(tau) return total_rewards[:, None] * total_scores def reinforce_variance(n_samples, T, n_feat, n_actions, seed=42): """ Computes the variance of the REINFORCE gradient estimator for trajectory length T. """ rng = np.random.default_rng(seed) all_g = simulate_reinforce_gradients(n_samples, T, n_feat, n_actions, rng) return np.var(all_g, axis=0).sum() def run_variance_experiment(T_list, n_samples, n_feat, n_actions): """ Runs the simulation over multiple trajectory lengths T. """ var_vs_T = [] for T in T_list: print(f"Running simulation for T={T}...") var_g = reinforce_variance(n_samples, T, n_feat, n_actions, seed=42) var_vs_T.append(var_g) return np.array(var_vs_T) def plot_variance_scaling(T_list, var_vs_T, output_path): """ Plots the variance scaling results in a 2-panel figure. """ log_T = np.log(T_list) log_var_T = np.log(var_vs_T) slope_T, intercept_T = np.polyfit(log_T, log_var_T, 1) print(f"\nEmpirical slope wrt T: {slope_T:.4f}") ratios = var_vs_T / (T_list ** 3) mean_ratio = np.mean(ratios) print(f"Mean normalized ratio Var(g)/T^3: {mean_ratio:.4f}") fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(8.5, 3.6)) for ax in (ax1, ax2): ax.set_facecolor("none") ax.spines["top"].set_visible(False) ax.spines["right"].set_visible(False) ax.spines["left"].set_color(SPINE_COLOR) ax.spines["bottom"].set_color(SPINE_COLOR) ax.tick_params(which='both', colors=COLOR_MUTED, labelsize=9) ax.grid(True, linestyle="--", alpha=0.4, color=GRID_COLOR) fig.patch.set_facecolor("none") # --- Panel 1: Var vs T --- ax1.loglog(T_list, var_vs_T, "o", color=COLOR_T, ms=6, label="Empirical variance", zorder=3) fit_T = np.exp(intercept_T) * (T_list ** slope_T) ax1.loglog(T_list, fit_T, "-", color=COLOR_T, lw=2, label=f"Fit: $\\propto T^{{{slope_T:.2f}}}$", zorder=2) # Reference line T^3 ref_T = var_vs_T[0] * (T_list / T_list[0])**3 ax1.loglog(T_list, ref_T, ":", color=COLOR_REF, lw=1.2, label=r"Ideal $O(T^3)$ slope", zorder=1) ax1.set_xlabel("Trajectory length $T$", fontsize=10, color=COLOR_TEXT) ax1.set_ylabel(r"Estimator variance $\mathrm{Var}(\hat{g})$", fontsize=10, color=COLOR_TEXT) ax1.set_title(r"Scaling with $T$", fontsize=11, fontweight="bold", color=COLOR_TEXT, pad=8) leg1 = ax1.legend(fontsize=8, loc="upper left", frameon=False) plt.setp(leg1.get_texts(), color=COLOR_TEXT) ax1.text(0.95, 0.08, f"Empirical slope: {slope_T:.2f}\n(Expected: 3.00)", transform=ax1.transAxes, fontsize=8.5, color=COLOR_TEXT, ha="right", bbox=dict(boxstyle="round,pad=0.3", facecolor="#1E222A", alpha=0.3, edgecolor="none")) # --- Panel 2: Normalized ratio Var(g) / T^3 --- x_pos = np.arange(len(T_list)) ax2.axhline(mean_ratio, color=COLOR_REF, lw=1.5, ls="--", label=f"Mean = {mean_ratio:.3f}", zorder=1) ax2.scatter(x_pos, ratios, color=COLOR_RATIO, s=35, zorder=3, label=r"$\mathrm{Var}(\hat{g})\,/\,T^3$") ax2.set_xticks(x_pos) ax2.set_xticklabels([f"T={t}" for t in T_list], rotation=30, ha="right", fontsize=8) ax2.set_ylim(mean_ratio * 0.7, mean_ratio * 1.3) ax2.set_xlabel("Trajectory length $T$", fontsize=10, color=COLOR_TEXT) ax2.set_ylabel(r"Normalized Variance Ratio", fontsize=10, color=COLOR_TEXT) ax2.set_title(r"Normalized $\frac{\mathrm{Var}(\hat{g})}{T^3}$", fontsize=11, fontweight="bold", color=COLOR_TEXT, pad=8) leg2 = ax2.legend(fontsize=8, loc="upper right", frameon=False) plt.setp(leg2.get_texts(), color=COLOR_TEXT) ax2.text(0.05, 0.08, f"Mean ratio: {mean_ratio:.3f}", transform=ax2.transAxes, fontsize=8.5, color=COLOR_TEXT, ha="left", bbox=dict(boxstyle="round,pad=0.3", facecolor="#1E222A", alpha=0.3, edgecolor="none")) plt.tight_layout() plt.savefig(output_path, dpi=200, bbox_inches="tight", transparent=True) print(f"\nSaved 2-panel figure to {output_path}") if __name__ == "__main__": n_samples = 30000 n_actions = 4 n_feat = 4 T_list = np.array([5, 10, 20, 40, 80, 160, 320, 640]) output_path = "/Users/pedregosa/dev/webpage/content/images/2026/reinforce_variance_scaling.png" var_vs_T = run_variance_experiment(T_list, n_samples, n_feat, n_actions) plot_variance_scaling(T_list, var_vs_T, output_path)