File size: 4,407 Bytes
510ab6f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
bf60b3f
 
a2fe35b
510ab6f
 
 
 
fe8455b
510ab6f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a2fe35b
510ab6f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a2fe35b
 
 
 
 
510ab6f
 
 
a2fe35b
 
 
 
 
 
 
510ab6f
a2fe35b
510ab6f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a2fe35b
 
 
510ab6f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a2fe35b
 
 
 
 
510ab6f
a2fe35b
510ab6f
 
 
a2fe35b
510ab6f
 
 
 
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
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
"""
FlexTime Inference Baseline Script
Required by Hackathon Spec
"""

import asyncio
import json
import os
import textwrap
from typing import List, Optional

# Ensure the root dir is in path
import sys
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from openai import OpenAI

from server.engine import FlexTimeEnv, TASK_CONFIGS
from server.models import Action
from agent.llm_agent import LLMAgent

# Mandatory environment variables with defaults
API_BASE_URL = os.getenv("API_BASE_URL", "https://api.openai.com/v1")
MODEL_NAME = os.getenv("MODEL_NAME", "gpt-4o-mini")
HF_TOKEN = os.getenv("HF_TOKEN")
BENCHMARK = "FlexTime"

MAX_STEPS = 120
TEMPERATURE = 0.0
MAX_TOKENS = 120


def log_start(task: str, env: str, model: str) -> None:
    print(f"[START] task={task} env={env} model={model}", flush=True)


def log_step(step: int, action: str, reward: float, done: bool, error: Optional[str]) -> None:
    error_val = error if error else "null"
    done_val = str(done).lower()
    print(
        f"[STEP] step={step} action={action} reward={reward:.2f} done={done_val} error={error_val}",
        flush=True,
    )


def log_end(success: bool, steps: int, score: float, rewards: List[float]) -> None:
    rewards_str = ",".join(f"{r:.2f}" for r in rewards)
    print(
        f"[END] success={str(success).lower()} steps={steps} score={score:.3f} rewards={rewards_str}",
        flush=True
    )


def run_task(agent: LLMAgent, env: FlexTimeEnv, task_id: str):
    log_start(task=task_id, env=BENCHMARK, model=MODEL_NAME)
    
    rewards: List[float] = []
    steps_taken = 0
    success = False
    score = 0.0
    
    try:
        # Reset the environment for the specific task using a consistent seed
        obs = env.reset(task_id=task_id, seed=42)
        obs_dict = obs.model_dump()
        done = False
        
        cfg = TASK_CONFIGS[task_id]
        cur_max_steps = min(MAX_STEPS, cfg["max_steps"])
        target_score = cfg["target_score"]
        
        agent.reset()
        
        last_reward = None
        last_error = None
        
        for step in range(1, cur_max_steps + 1):
            if done:
                break
                
            # Smart Early Termination
            if len(rewards) >= 3 and rewards[-1] == rewards[-2] == rewards[-3]:
                # Terminate if same reward repeats 3 times (no progress)
                done = True
                break
                
            # Predict
            action_dict = agent.generate_action(obs_dict, last_reward, last_error)
            action_str = json.dumps(action_dict).replace(' ', '')
            
            # Execute
            try:
                action = Action(**action_dict)
                result = env.step(action)
                obs_dict = result.observation.model_dump()
                reward = result.reward.total or 0.0
                done = result.done
                error = None
            except Exception as e:
                # Execution failed due to malformed action
                reward = -0.05
                done = False
                error = str(e).replace(' ', '_') # Replace spaces just in case format is very strictly space-delimited
            
            rewards.append(reward)
            steps_taken = step
            
            last_reward = reward
            last_error = error
            
            log_step(step=step, action=action_str, reward=reward, done=done, error=error)
        
        # Grading
        grade = env.grade()
        score = grade.get("score", 0.0)
        score = min(max(score, 0.0), 1.0) # Clamp 0-1
        success = score >= target_score
        
    except Exception as overall_e:
        print(f"[DEBUG] Overall environment failure: {overall_e}", flush=True)
    finally:
        log_end(success=success, steps=steps_taken, score=score, rewards=rewards)


def main():
    if not HF_TOKEN:
        print("[DEBUG] HF_TOKEN is missing. This will crash. Please set HF_TOKEN.", flush=True)
        # We allow client initialization crash if token is missing as it enforces the constraint.
        
    client = OpenAI(base_url=API_BASE_URL, api_key=HF_TOKEN or "dummy-key")
    env = FlexTimeEnv()
    agent = LLMAgent(client, MODEL_NAME)
    
    tasks = ["task_easy", "task_medium", "task_hard"]
    for t_id in tasks:
        run_task(agent, env, t_id)


if __name__ == "__main__":
    main()