""" UGTC: Uncertainty-Gated Temporal Credit — Interactive Demo =========================================================== This Gradio app provides an educational interface for understanding the UGTC algorithm, its mathematics, and its RL integrations. Paper: https://doi.org/10.5281/zenodo.19715116 GitHub: https://github.com/ethosoftai/ugtc """ import numpy as np import torch import torch.nn as nn import gradio as gr import plotly.graph_objects as go import plotly.express as px from plotly.subplots import make_subplots # ── UGTC Core (self-contained, no external dep) ─────────────────────────────── def sigmoid(x): return 1.0 / (1.0 + np.exp(-x)) def compute_gate(sigma, sigma_ema, beta=5.0, eps=1e-8): """Compute UGTC uncertainty gate u(s) = sigmoid(-β·(σ̂ - 1)).""" sigma_hat = sigma / (sigma_ema + eps) return sigmoid(-beta * (sigma_hat - 1.0)) def compute_gae(rewards, values, next_values, dones, gamma=0.99, lam=0.9): """Standard GAE computation.""" T = len(rewards) advantages = np.zeros(T) gae = 0.0 for t in reversed(range(T)): delta = rewards[t] + gamma * next_values[t] * (1 - dones[t]) - values[t] gae = delta + gamma * lam * (1 - dones[t]) * gae advantages[t] = gae return advantages def blend_advantages(adv_fast, adv_slow, gates): """A^UGTC = u(s)·A^slow + (1-u(s))·A^fast""" return gates * adv_slow + (1 - gates) * adv_fast # ── Tab 1: Gate Visualizer ──────────────────────────────────────────────────── def plot_gate_curve(beta, sigma_ema): """Plot gate u(s) as a function of ensemble disagreement σ(s).""" sigma_vals = np.linspace(0, 3 * sigma_ema, 500) gate_vals = compute_gate(sigma_vals, sigma_ema, beta=beta) fig = go.Figure() fig.add_trace(go.Scatter( x=sigma_vals, y=gate_vals, mode='lines', name='u(σ)', line=dict(color='#6366f1', width=3), )) fig.add_hline(y=0.5, line_dash='dot', line_color='#94a3b8', annotation_text='u = 0.5 (equal blend)') fig.add_vline(x=float(sigma_ema), line_dash='dash', line_color='#10b981', annotation_text=f'σ_EMA = {sigma_ema:.2f}') # Shade regions fig.add_vrect(x0=0, x1=float(sigma_ema), fillcolor='rgba(16,185,129,0.08)', line_width=0, annotation_text='Use A^slow', annotation_position='top left') fig.add_vrect(x0=float(sigma_ema), x1=float(3 * sigma_ema), fillcolor='rgba(239,68,68,0.06)', line_width=0, annotation_text='Use A^fast', annotation_position='top right') fig.update_layout( title=f'UGTC Gate u(σ) — β={beta:.1f}, σ_EMA={sigma_ema:.2f}', xaxis_title='Ensemble Disagreement σ(s)', yaxis_title='Gate Value u(s)', yaxis=dict(range=[-0.05, 1.05]), template='plotly_dark', height=420, plot_bgcolor='rgba(17,24,39,1)', paper_bgcolor='rgba(17,24,39,1)', font=dict(color='#f1f5f9'), ) return fig def gate_demo(beta, sigma_ema, current_sigma): """Show current gate value for a specific σ(s) point.""" gate = float(compute_gate(current_sigma, sigma_ema, beta=beta)) fig = plot_gate_curve(beta, sigma_ema) # Mark the current point fig.add_trace(go.Scatter( x=[current_sigma], y=[gate], mode='markers', name='Current state', marker=dict(color='#f59e0b', size=14, symbol='star'), )) blend_desc = "slow (accurate)" if gate > 0.6 else ("fast (stable)" if gate < 0.4 else "equal blend") action = f"→ {'more' if gate > 0.6 else 'less'} weight on slow estimate" summary = f""" **Current state:** - σ(s) = {current_sigma:.3f} | σ_EMA = {sigma_ema:.3f} | σ̂ = {current_sigma/sigma_ema:.3f} - Gate u(s) = **{gate:.4f}** - Blending: {gate:.1%} slow + {1-gate:.1%} fast - Decision: use **{blend_desc}** estimate """ return fig, summary # ── Tab 2: Advantage Blending ───────────────────────────────────────────────── def plot_advantage_blend(lambda_fast, lambda_slow, beta, seed): """Simulate UGTC advantage blending over a trajectory.""" rng = np.random.default_rng(int(seed)) T = 64 gamma = 0.99 rewards = rng.normal(0, 1, T) v_fast = rng.normal(0, 0.5, T) v_slow1 = v_fast + rng.normal(0, 0.3, T) v_slow2 = v_fast + rng.normal(0, 0.3, T) v_slow3 = v_fast + rng.normal(0, 0.3, T) v_slow_mean = (v_slow1 + v_slow2 + v_slow3) / 3 sigma = np.std([v_slow1, v_slow2, v_slow3], axis=0) nv_fast = np.roll(v_fast, -1); nv_fast[-1] = 0 nv_slow = np.roll(v_slow_mean, -1); nv_slow[-1] = 0 dones = np.zeros(T); dones[-1] = 1.0 adv_fast = compute_gae(rewards, v_fast, nv_fast, dones, gamma, lambda_fast) adv_slow = compute_gae(rewards, v_slow_mean, nv_slow, dones, gamma, lambda_slow) sigma_ema = float(sigma.mean()) gates = compute_gate(sigma, sigma_ema, beta=beta) adv_ugtc = blend_advantages(adv_fast, adv_slow, gates) t = np.arange(T) fig = make_subplots( rows=3, cols=1, subplot_titles=[ 'Advantages (Fast vs Slow vs UGTC Blended)', 'Uncertainty Gate u(s) over Trajectory', 'Ensemble Disagreement σ(s)', ], vertical_spacing=0.1, ) fig.add_trace(go.Scatter(x=t, y=adv_fast, name='A^fast (λ={:.2f})'.format(lambda_fast), line=dict(color='#f59e0b', width=1.5, dash='dot')), row=1, col=1) fig.add_trace(go.Scatter(x=t, y=adv_slow, name='A^slow (λ={:.2f})'.format(lambda_slow), line=dict(color='#06b6d4', width=1.5, dash='dash')), row=1, col=1) fig.add_trace(go.Scatter(x=t, y=adv_ugtc, name='A^UGTC (blended)', line=dict(color='#6366f1', width=2.5)), row=1, col=1) fig.add_trace(go.Scatter(x=t, y=gates, name='Gate u(s)', fill='tozeroy', line=dict(color='#10b981'), fillcolor='rgba(16,185,129,0.15)'), row=2, col=1) fig.add_trace(go.Bar(x=t, y=sigma, name='σ(s)', marker_color='#8b5cf6', opacity=0.7), row=3, col=1) fig.update_layout( template='plotly_dark', height=700, plot_bgcolor='rgba(17,24,39,1)', paper_bgcolor='rgba(17,24,39,1)', font=dict(color='#f1f5f9'), showlegend=True, ) fig.update_xaxes(title_text='Timestep') stats = f""" **Trajectory Statistics (T={T}):** - A^fast: mean={adv_fast.mean():.3f}, std={adv_fast.std():.3f} - A^slow: mean={adv_slow.mean():.3f}, std={adv_slow.std():.3f} - A^UGTC: mean={adv_ugtc.mean():.3f}, std={adv_ugtc.std():.3f} - Mean gate u(s): {gates.mean():.3f} → {gates.mean()*100:.0f}% slow weight on average - σ_EMA: {sigma_ema:.4f} """ return fig, stats # ── Tab 3: Algorithm Comparison ─────────────────────────────────────────────── ALGO_CONTENT = { "UGTC-PPO": { "description": """ ## UGTC-PPO Integration **Key change from vanilla PPO:** Advantages are computed by `UGTCModule.compute_advantages()` instead of standard single-critic GAE. **Actor loss (unchanged PPO clipped surrogate):** ``` ratio = π_θ(a|s) / π_θ_old(a|s) L_policy = -mean(min(ratio · Â^UGTC, clip(ratio, 1±ε) · Â^UGTC)) ``` **Critic losses (train all UGTC critics):** ``` L_fast = MSE(V_fast(s), R̂) L_slow = mean over m of MSE(Vᵐ_slow(s), R̂) L_total = L_policy + 0.5·(L_fast + L_slow) ``` Where `R̂ = Â^UGTC + V^UGTC(s)` are the value targets. **What stays the same:** PPO clip coefficient ε, entropy bonus, gradient clipping, rollout collection — all unchanged. **Overhead:** 3 extra MLP value heads (fast + 3 slow). Negligible relative to actor network. """, "pseudocode": """ REPEAT until convergence: τ ← collect_rollout(π_θ) # same as vanilla PPO # === KEY CHANGE: UGTC advantage === A^UGTC ← ugtc.compute_advantages(τ, γ) Â ← normalize(A^UGTC) R̂ ← Â + V^UGTC(s) FOR k = 1 TO K: # standard PPO epochs ratio ← π_θ(a|s) / π_old(a|s) L_policy ← clipped_surrogate(ratio, Â) L_critics ← MSE(V_fast, R̂) + MSE(V_slow, R̂) step(L_policy + 0.5·L_critics) """ }, "UGTC-TD3": { "description": """ ## UGTC-TD3 Integration **Key change from vanilla TD3:** The actor receives a UGTC baseline correction. **Actor loss (modified):** ``` Q_min = min(Q¹(s, π(s)), Q²(s, π(s))) V^UGTC(s) = u(s)·V̄_slow(s) + (1-u(s))·V_fast(s) A^UGTC(s) = Q_min - V^UGTC(s) L_actor = -mean(Q_min + η·A^UGTC(s)) ``` **Critic loss (unchanged standard TD3 twin-Q):** ``` ỹ = r + γ·min(Q¹_target, Q²_target)(s', π_target(s') + ε) L_critic = MSE(Q¹(s,a), ỹ) + MSE(Q²(s,a), ỹ) ``` **What stays the same:** Twin-Q, target policy smoothing, delayed actor update, soft target updates — all unchanged. **Note:** η=0.5 (correction weight) is an implementation default, not a fixed UGTC hyperparameter. May benefit from tuning. """, "pseudocode": """ REPEAT: B = sample(replay_buffer) # Critic update (standard TD3 — unchanged) ỹ = r + γ·min(Q_target)(s', π_target(s') + ε) step(MSE(Q¹(s,a), ỹ) + MSE(Q²(s,a), ỹ)) IF step % policy_delay == 0: # === KEY CHANGE: UGTC baseline correction === Q_min = min(Q¹(s, π(s)), Q²(s, π(s))) V_ugtc = u(s)·V̄_slow + (1-u(s))·V_fast A_ugtc = Q_min - V_ugtc step(-mean(Q_min + η·A_ugtc)) # actor loss soft_update(targets) """ }, "UGTC-SAC": { "description": """ ## UGTC-SAC Integration **Key change from vanilla SAC:** V^UGTC replaces the implicit value baseline in the actor loss. **Standard SAC actor loss:** ``` L_π = mean(α·log π(a|s) - Q_min(s, a)) ``` **UGTC-SAC actor loss:** ``` V^UGTC(s) = u(s)·V̄_slow(s) + (1-u(s))·V_fast(s) L_π^UGTC = mean(α·log π(a|s) - Q_min(s, a) + V^UGTC(s)) ``` By adding V^UGTC(s), the actor's gradient is centered around a more accurate baseline, reducing variance. **What stays the same:** Entropy coefficient α, twin-Q critic training, automatic entropy tuning, target network updates — all unchanged. """, "pseudocode": """ REPEAT: B = sample(replay_buffer) # Critic update (standard SAC — unchanged) ã' = sample π(·|s') y = r + γ(Q_min_target(s',ã') - α·log π(ã'|s')) step(MSE(Q¹(s,a), y) + MSE(Q²(s,a), y)) # === KEY CHANGE: V^UGTC as actor baseline === ã, log_π = sample π(·|s) Q_min = min(Q¹(s,ã), Q²(s,ã)) V_ugtc = u(s)·V̄_slow + (1-u(s))·V_fast # Standard: mean(α·log π - Q_min) # UGTC-SAC: mean(α·log π - Q_min + V_ugtc) step(mean(α·log_π - Q_min + V_ugtc)) """ }, } def show_algorithm(algo_name): content = ALGO_CONTENT[algo_name] return content["description"], f"```\n{content['pseudocode']}\n```" # ── Tab 4: Hyperparameter Explorer ──────────────────────────────────────────── def plot_beta_effect(beta_values_str): """Show how different β values affect gate sharpness.""" sigma_range = np.linspace(0, 3, 400) sigma_ema = 1.0 beta_list = [float(b.strip()) for b in beta_values_str.split(',') if b.strip()] colors = ['#6366f1', '#10b981', '#f59e0b', '#ef4444', '#06b6d4'] fig = go.Figure() for i, beta in enumerate(beta_list[:5]): gates = compute_gate(sigma_range, sigma_ema, beta=beta) fig.add_trace(go.Scatter( x=sigma_range, y=gates, name=f'β = {beta}', line=dict(color=colors[i % len(colors)], width=2.5), )) fig.add_vline(x=1.0, line_dash='dash', line_color='#94a3b8', annotation_text='σ_EMA = 1.0') fig.add_hline(y=0.5, line_dash='dot', line_color='#94a3b8') fig.update_layout( title='Gate Sharpness: Effect of Temperature β', xaxis_title='Normalized Disagreement σ̂(s) = σ(s)/σ_EMA', yaxis_title='Gate Value u(s)', yaxis=dict(range=[-0.05, 1.05]), template='plotly_dark', height=420, plot_bgcolor='rgba(17,24,39,1)', paper_bgcolor='rgba(17,24,39,1)', font=dict(color='#f1f5f9'), ) desc = ( "**β interpretation:**\n" "- Low β (e.g. 0.5): Gradual, smooth blending\n" "- β = 5.0 (paper default): Moderate sharpness, responsive to uncertainty\n" "- High β (e.g. 20): Near-binary switching between fast and slow\n\n" "The paper uses β = 5.0 across all benchmarks without tuning." ) return fig, desc def plot_lambda_effect(): """Show how λ_fast vs λ_slow affects GAE horizon.""" T = 50 gamma = 0.99 t = np.arange(T) fig = go.Figure() for lam, name, color in [ (0.80, 'λ_fast=0.80 (low bias-var tradeoff)', '#f59e0b'), (0.95, 'λ=0.95 (standard GAE)', '#94a3b8'), (0.99, 'λ_slow=0.99 (long horizon)', '#6366f1'), ]: # Effective weight of step k in GAE: (γλ)^k weights = np.array([(gamma * lam) ** k for k in range(T)]) weights /= weights.sum() fig.add_trace(go.Scatter(x=t, y=weights, name=name, line=dict(color=color, width=2.5))) fig.update_layout( title='GAE Credit Attribution Weights by λ Value', xaxis_title='Steps into future (k)', yaxis_title='Relative weight of δₜ₊ₖ in Aₜ', template='plotly_dark', height=400, plot_bgcolor='rgba(17,24,39,1)', paper_bgcolor='rgba(17,24,39,1)', font=dict(color='#f1f5f9'), ) return fig # ── Tab 5: Math ─────────────────────────────────────────────────────────────── MATH_SECTIONS = { "GAE Foundation": """ ## Generalized Advantage Estimation The TD residual at step t: **δₜ = rₜ + γ·V(sₜ₊₁)·(1 - dₜ) - V(sₜ)** The GAE advantage (reverse accumulation): **Aₜ^GAE(λ) = Σₖ₌₀^∞ (γλ)ᵏ · δₜ₊ₖ** This is a weighted sum of k-step TD residuals. The λ parameter controls the effective horizon: - **λ → 0**: Pure TD(0) — low variance, high bias (myopic) - **λ → 1**: Monte Carlo — low bias, high variance (long-sighted) UGTC exploits this by maintaining two separate GAE streams with different λ values. """, "Dual Critics": """ ## UGTC Dual-Critic Architecture **Fast critic** (single MLP): ``` V_fast: obs_dim → 64 → 64 → 1 A^fast_t = GAE(τ, V=V_fast, λ=λ_fast=0.80) ``` Characteristic: Low variance (single point estimate), higher bias **Slow ensemble** (M=3 independent MLPs): ``` V¹_slow, V²_slow, V³_slow (same arch, different random init) V̄_slow = mean(V¹, V², V³) A^slow_t = GAE(τ, V=V̄_slow, λ=λ_slow=0.99) ``` Characteristic: Lower bias (long horizon + ensemble mean), higher variance **Key insight:** The two streams provide complementary error profiles. The gate selects between them based on per-state confidence. """, "Uncertainty Gate": """ ## The Uncertainty Gate **Step 1** — Measure ensemble disagreement: ``` σ(s) = std(V¹_slow(s), V²_slow(s), V³_slow(s)) ``` High σ means the ensemble members disagree → the slow estimate is unreliable. **Step 2** — EMA normalization (momentum α = 0.99): ``` σ_EMA ← α · σ_EMA + (1-α) · E[σ(s)] σ̂(s) = σ(s) / (σ_EMA + ε) ``` This converts absolute disagreement to a relative measure. σ̂ > 1 means "above-average uncertainty." **Step 3** — Sigmoid gate (temperature β = 5.0): ``` u(s) = sigmoid(-β · (σ̂(s) - 1.0)) ``` When σ̂ = 1 (average uncertainty): u = 0.5 (equal blend) When σ̂ < 1 (below average): u → 1 (use slow) When σ̂ > 1 (above average): u → 0 (use fast) """, "Blended Advantage": """ ## UGTC Blended Advantage The final blended advantage: **A^UGTC_t = u(sₜ) · A^slow_t + (1 - u(sₜ)) · A^fast_t** And the blended value: **V^UGTC(s) = u(s) · V̄_slow(s) + (1-u(s)) · V_fast(s)** **Intuition:** - State well-understood (ensemble agrees, σ̂ < 1) → u(s) → 1 → use long-horizon A^slow - State uncertain (ensemble disagrees, σ̂ > 1) → u(s) → 0 → use stable A^fast This is analogous to Bayesian model averaging, but with a deterministic gate based on functional disagreement rather than posterior uncertainty. **In PPO:** `A^UGTC` directly replaces the GAE advantage in the clipped surrogate. **In TD3/SAC:** `V^UGTC` replaces the value baseline in the actor gradient. """, } def show_math(section): return MATH_SECTIONS[section] # ── Layout ──────────────────────────────────────────────────────────────────── HEADER = """
A plug-in advantage estimator for actor-critic reinforcement learning
Accepted · Ulysseus Young Explorers in Science (UYES) Journal · Journal DOI forthcoming