Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import plotly.graph_objects as go | |
| import pandas as pd | |
| import numpy as np | |
| import torch | |
| import os | |
| import sys | |
| from fastapi import FastAPI, HTTPException | |
| # Ensure local packages (bess_rl subset) are importable | |
| sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "."))) | |
| from openenv.server.env import BESSEnvironment | |
| from agent.actor_critic import SAC_Agent | |
| from agent.config import AgentConfig | |
| from openenv.models import ActionModel, ObservationModel, StepResult, ResetConfig | |
| # --- Configuration & Styling --- | |
| CUSTOM_CSS = """ | |
| @import url('https://fonts.googleapis.com/css2?family=Orbitron:wght@400;700&display=swap'); | |
| .gradio-container { | |
| background: radial-gradient(circle at top right, #0d1117 0%, #010409 100%) !important; | |
| font-family: 'Inter', system-ui, -apple-system, sans-serif !important; | |
| } | |
| h1, h2, h3 { | |
| font-family: 'Orbitron', sans-serif !important; | |
| text-transform: uppercase; | |
| letter-spacing: 2px; | |
| color: #58a6ff !important; | |
| } | |
| .kpi-card { | |
| background: rgba(22, 27, 34, 0.6) !important; | |
| border: 1px solid rgba(88, 166, 255, 0.2) !important; | |
| border-radius: 12px !important; | |
| padding: 20px !important; | |
| box-shadow: 0 4px 30px rgba(0, 0, 0, 0.1); | |
| backdrop-filter: blur(5px); | |
| transition: transform 0.3s ease, border-color 0.3s ease; | |
| } | |
| .kpi-card:hover { | |
| transform: translateY(-5px); | |
| border-color: rgba(88, 166, 255, 0.6) !important; | |
| } | |
| .neon-text { | |
| text-shadow: 0 0 10px rgba(88, 166, 255, 0.5); | |
| } | |
| .sidebar { | |
| background: transparent !important; | |
| } | |
| """ | |
| def load_agent(difficulty): | |
| config = AgentConfig() | |
| agent = SAC_Agent(config) | |
| # Path to models in the repo | |
| model_dir = os.path.join(os.path.dirname(__file__), "train", "models") | |
| model_name = f"best_model_{difficulty}" | |
| model_path = os.path.join(model_dir, model_name) | |
| if os.path.exists(model_path + "_actor.pth"): | |
| try: | |
| # We only need the actor for inference | |
| agent.actor.load_state_dict(torch.load(model_path + "_actor.pth", map_location=torch.device('cpu'), weights_only=True)) | |
| print(f"Successfully loaded {difficulty} Soft Actor-Critic model.") | |
| except Exception as e: | |
| print(f"Error loading model: {e}") | |
| else: | |
| print(f"Model {model_name} not found. Using initialized random weights.") | |
| return agent | |
| def run_simulation(difficulty, capacity, max_power, episode_length): | |
| # Set the path safely relative to the space execution | |
| env = BESSEnvironment(data_path=os.path.join(os.path.dirname(__file__), "data", "pjm_data.csv")) | |
| env.reset(seed=42, task=difficulty) | |
| # Update env config | |
| env.config.capacity_mwh = capacity | |
| env.config.max_charge_mw = max_power | |
| env.config.max_discharge_mw = max_power | |
| agent = load_agent(difficulty) | |
| history = [] | |
| # Seed 42 is historically interesting for the visual output | |
| state = env.reset(seed=42, task=difficulty) | |
| total_reward = 0 | |
| total_ea_profit = 0 | |
| total_fr_reward = 0 | |
| peak_reduction = 0 | |
| steps = min(int(episode_length), env.max_steps) | |
| for _ in range(steps): | |
| # Format state for agent | |
| state_arr = np.array([ | |
| state.hour_of_day, state.soc, state.price_lmp, | |
| state.p_avg, state.freq_regd, state.load_mw | |
| ]) | |
| # Select deterministic action using SAC inference profile | |
| action_vals = agent.select_action(state_arr, evaluate=True) | |
| action_model = ActionModel(action=action_vals.tolist()) | |
| step_result = env.step(action_model) | |
| # Track data | |
| res = step_result.info | |
| res['step'] = _ | |
| res['reward'] = step_result.reward | |
| history.append(res) | |
| total_reward += step_result.reward | |
| total_ea_profit += res.get('r_ea', 0) | |
| total_fr_reward += res.get('r_fr', 0) | |
| # Baseline vs Net Load for peak reduction | |
| if res.get('baseline_load', 0) > 20: | |
| reduction = res['baseline_load'] - res['net_load'] | |
| peak_reduction += max(0, reduction) | |
| state = step_result.observation | |
| if step_result.terminated: | |
| break | |
| df = pd.DataFrame(history) | |
| # Summary Metrics | |
| roi = (total_reward / (capacity * 1000)) * 100 # Rough ROI estimate | |
| roi_str = f"{roi:.2f}%" | |
| total_profit_str = f"${total_reward:,.0f}" | |
| # Grid Stability: Function of successful FR accuracy and lack of load violations | |
| fr_hits = len(df[df['r_fr'] > 0]) if 'r_fr' in df else 0 | |
| stab_score = min(99.9, max(50.0, 70 + (fr_hits / max(1, steps)) * 30)) | |
| grid_stability = f"{stab_score:.1f}%" | |
| return df, total_profit_str, roi_str, grid_stability | |
| def create_plots(df): | |
| # 1. Price vs SOC | |
| fig1 = go.Figure() | |
| fig1.add_trace(go.Scatter(x=df['step'], y=df['lmp'], name='LMP ($/MWh)', line=dict(color='#f1c40f', width=2))) | |
| fig1.add_trace(go.Scatter(x=df['step'], y=df['soc']*100, name='SOC (%)', line=dict(color='#2ecc71', width=3), yaxis='y2')) | |
| # Fill background when SOC is extremely low to indicate warning | |
| for index, row in df.iterrows(): | |
| if row['soc'] < 0.2: | |
| fig1.add_vrect(x0=index-0.5, x1=index+0.5, fillcolor="rgba(231, 76, 60, 0.1)", layer="below", line_width=0) | |
| fig1.update_layout( | |
| template='plotly_dark', | |
| title='Locational Marginal Price vs Battery State of Charge', | |
| xaxis_title='Step (Hour)', | |
| yaxis=dict(title='Price ($/MWh)', side='left'), | |
| yaxis2=dict(title='SOC (%)', overlaying='y', side='right', range=[0, 105]), | |
| margin=dict(l=20, r=20, t=40, b=20), | |
| legend=dict(orientation="h", yanchor="bottom", y=1.02, xanchor="right", x=1) | |
| ) | |
| # 2. Net Load vs Baseline | |
| fig2 = go.Figure() | |
| fig2.add_trace(go.Scatter(x=df['step'], y=df['baseline_load'], name='Baseline Load', line=dict(color='#e74c3c', dash='dash'))) | |
| fig2.add_trace(go.Scatter(x=df['step'], y=df['net_load'], name='Grid-Aware Net Load', fill='tozeroy', line=dict(color='#58a6ff'))) | |
| fig2.update_layout( | |
| template='plotly_dark', | |
| title='Peak Shaving Performance', | |
| xaxis_title='Step (Hour)', | |
| yaxis_title='Load (MW)', | |
| margin=dict(l=20, r=20, t=40, b=20), | |
| legend=dict(orientation="h", yanchor="bottom", y=1.02, xanchor="right", x=1) | |
| ) | |
| # 3. Action Breakdown | |
| fig3 = go.Figure() | |
| fig3.add_trace(go.Bar(x=df['step'], y=df['action_ea'], name='Arbitrage', marker_color='#2980b9')) | |
| fig3.add_trace(go.Bar(x=df['step'], y=df['action_fr'], name='Frequency Reg', marker_color='#9b59b6')) | |
| fig3.add_trace(go.Bar(x=df['step'], y=df['action_ps'], name='Peak Shaving', marker_color='#e67e22')) | |
| fig3.update_layout( | |
| template='plotly_dark', | |
| title='Multi-Objective Control Actions', | |
| barmode='relative', | |
| xaxis_title='Step (Hour)', | |
| yaxis_title='Action Intensity (-1 to 1)', | |
| margin=dict(l=20, r=20, t=40, b=20), | |
| legend=dict(orientation="h", yanchor="bottom", y=1.02, xanchor="right", x=1) | |
| ) | |
| return fig1, fig2, fig3 | |
| def update_ui(difficulty, capacity, max_power, episode_length): | |
| df, prof, roi, stab = run_simulation(difficulty, capacity, max_power, episode_length) | |
| p1, p2, p3 = create_plots(df) | |
| return p1, p2, p3, prof, roi, stab | |
| with gr.Blocks(css=CUSTOM_CSS, theme=gr.themes.Monochrome()) as demo: | |
| gr.HTML("<h1 style='text-align: center; margin-bottom: 20px;' class='neon-text'>⚡ EnergyStock: RL Battery Terminal</h1>") | |
| gr.HTML("<p style='text-align: center; color: #8b949e;'>Co-Optimizing Energy Arbitrage, Frequency Regulation, and Peak Shaving using Soft Actor-Critic (SAC).</p>") | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| with gr.Group(): | |
| gr.HTML("<h3>⚙️ Control Panel</h3>") | |
| difficulty = gr.Radio(["easy", "medium", "hard"], label="Complexity Level", value="hard") | |
| capacity = gr.Slider(50, 500, value=150, label="Battery Capacity (MWh)") | |
| max_power = gr.Slider(10, 100, value=50, label="Max Power (MW)") | |
| episode_length = gr.Slider(24, 720, value=168, label="Simulation Hours") # Default to 1 week | |
| btn = gr.Button("🚀 Run Terminal Simulation", variant="primary") | |
| with gr.Column(scale=3): | |
| with gr.Row(): | |
| with gr.Column(elem_classes="kpi-card"): | |
| gr.HTML("<small>Estimated Profit</small>") | |
| profit_val = gr.HTML("<h2 class='neon-text'>$0</h2>") | |
| with gr.Column(elem_classes="kpi-card"): | |
| gr.HTML("<small>Projected ROI</small>") | |
| roi_val = gr.HTML("<h2 class='neon-text'>0%</h2>") | |
| with gr.Column(elem_classes="kpi-card"): | |
| gr.HTML("<small>Grid Stability Index</small>") | |
| stab_val = gr.HTML("<h2 class='neon-text'>0%</h2>") | |
| with gr.Tabs(): | |
| with gr.Tab("Market Analysis"): | |
| plot_market = gr.Plot() | |
| with gr.Tab("Grid Impact"): | |
| plot_grid = gr.Plot() | |
| with gr.Tab("Agent Actions (Nerd Mode)"): | |
| plot_nerd = gr.Plot() | |
| gr.Markdown("### Technical Details: Deep Reinforcement Learning Architecture") | |
| gr.Markdown(""" | |
| The agent utilizes a **Soft Actor-Critic (SAC)** algorithm capable of extreme exploration and exploitation, with dynamic entropy temperature scaling. | |
| - **Objective 1 (EA):** Buy low, sell high based on diurnal LMP patterns. | |
| - **Objective 2 (FR):** Minimize signal tracking error continuously against rapid regulation deployment frequencies. | |
| - **Objective 3 (PS):** Hoard physical SOC to deploy during Grid Baseline events breaking 20MW to alleviate stress. | |
| """) | |
| btn.click( | |
| update_ui, | |
| inputs=[difficulty, capacity, max_power, episode_length], | |
| outputs=[plot_market, plot_grid, plot_nerd, profit_val, roi_val, stab_val] | |
| ) | |
| # Run once on load | |
| demo.load( | |
| update_ui, | |
| inputs=[difficulty, capacity, max_power, episode_length], | |
| outputs=[plot_market, plot_grid, plot_nerd, profit_val, roi_val, stab_val] | |
| ) | |
| # --------------------------------------------------------------------------- | |
| # FastAPI application | |
| # --------------------------------------------------------------------------- | |
| app = FastAPI( | |
| title="EnergyStock: BESS-RL API", | |
| version="1.0.0", | |
| description="REST API for the SAC Battery Energy Storage RL agent + Gradio UI" | |
| ) | |
| # Shared environment singleton (stateful, one active session at a time) | |
| _DATA_PATH = os.path.join(os.path.dirname(__file__), "data", "pjm_data.csv") | |
| _env = BESSEnvironment(data_path=_DATA_PATH) | |
| def health(): | |
| """Liveness probe - confirms the server is running.""" | |
| return {"status": "ok", "message": "BESS-RL API is healthy"} | |
| def api_reset(config: ResetConfig): | |
| """Reset the simulation environment and return the initial observation.""" | |
| try: | |
| obs = _env.reset(seed=config.seed, task=config.task) | |
| return obs | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=str(e)) | |
| def api_step(action: ActionModel): | |
| """Advance the simulation by one timestep using the provided action vector.""" | |
| try: | |
| result = _env.step(action) | |
| return result | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=str(e)) | |
| def api_state(): | |
| """Return the current environment observation without advancing the simulation.""" | |
| try: | |
| return _env._get_obs() | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=str(e)) | |
| def api_info(): | |
| """Return metadata about the current simulation session.""" | |
| return { | |
| "task": _env.task, | |
| "max_steps": _env.max_steps, | |
| "current_step": _env.current_step, | |
| "soc": _env.soc, | |
| } | |
| # Mount the Gradio UI at the root - all /api/* routes are defined above | |
| # so they take precedence over the Gradio catch-all. | |
| app = gr.mount_gradio_app(app, demo, path="/") | |