File size: 2,425 Bytes
88bc772 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 | import os
import json
import numpy as np
from dotenv import load_dotenv
import google.generativeai as genai
from server.battery_environment import BatteryEnvironment
from models import BatteryAction
# Load .env
load_dotenv()
API_KEY = os.getenv("GEMINI_API_KEY")
def grade_trajectory(total_reward, mode):
max_theoretical = {"easy": 5000.0, "medium": 2000.0, "hard": 500.0}
target = max_theoretical.get(mode, 1000.0)
score = max(0.0, min(1.0, total_reward / target))
return score
def run_baseline(mode, model_name="gemini-1.5-flash"):
if not API_KEY:
print("GEMINI_API_KEY not found in .env. Falling back to Rule-Based logic.")
return 0.0
genai.configure(api_key=API_KEY)
model = genai.GenerativeModel(model_name)
env = BatteryEnvironment(mode=mode)
obs = env.reset()
total_reward = 0.0
done = False
step_count = 0
max_steps = 96
print(f"\n--- Starting Task: {mode.upper()} ---")
while not done and step_count < max_steps:
prompt = (f"You are a battery management agent. Maximize profit. "
f"Return action as JSON: {{'market_choice': [0.0-1.0], 'p_fraction': [-1.0-1.0]}}. "
f"Current State: energy={obs.energy_price}, fcl={obs.fcr_price}, soc={obs.soc}")
try:
response = model.generate_content(prompt, generation_config={"response_mime_type": "application/json"})
action_json = json.loads(response.text)
action = BatteryAction(
market_choice=float(action_json.get("market_choice", 0.0)),
p_fraction=float(action_json.get("p_fraction", 0.0))
)
except Exception as e:
action = BatteryAction(market_choice=0.0, p_fraction=0.0)
obs = env.step(action)
total_reward += obs.reward or 0.0
done, step_count = obs.done, step_count + 1
score = grade_trajectory(total_reward, mode)
print(f"Task '{mode}' completed. Total Reward: {total_reward:.2f}. Score: {score:.2f} / 1.00")
return score
if __name__ == "__main__":
scores = {}
for task_mode in ["easy", "medium", "hard"]:
score = run_baseline(task_mode)
scores[task_mode] = score
print("\n=== Final Gemini Baseline Scores ===")
for task_mode, score in scores.items():
print(f"{task_mode.capitalize()}: {score:.2f}")
|