Spaces:
Sleeping
Sleeping
File size: 5,734 Bytes
75c7554 | 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 142 143 144 145 | import os
import sys
import numpy as np
import time
import subprocess
import argparse
import json
# Ensure project root is on sys.path
_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
if _ROOT not in sys.path:
sys.path.insert(0, _ROOT)
# Canonical models directory (matches llm_evaluate.py and evaluate.py)
_MODELS_DIR = os.path.join(os.path.dirname(__file__), 'models')
from openenv.client import OpenEnvClient
from agent.config import AgentConfig
from agent.replay_buffer import ReplayBuffer
from agent.actor_critic import SAC_Agent
def start_server():
import requests
# Check if a server is already running (e.g. from Docker)
try:
r = requests.get("http://127.0.0.1:8000/api/health", timeout=1)
if r.status_code == 200:
print("Detected existing OpenEnv Server (likely Docker). Using it.")
return None
except:
pass
print("Starting local OpenEnv Server on port 8000...")
env_vars = os.environ.copy()
env_vars["PYTHONPATH"] = _ROOT
server_process = subprocess.Popen(
[sys.executable, "backend/main.py"], # Use our established main.py
env=env_vars,
stdout=subprocess.DEVNULL,
stderr=subprocess.DEVNULL
)
time.sleep(5) # Wait for server startup
return server_process
def train(args):
server_process = start_server()
try:
config = AgentConfig()
client = OpenEnvClient(base_url="http://127.0.0.1:8000")
agent = SAC_Agent(config)
load_path = args.load_model if args.load_model else os.path.join(_MODELS_DIR, f"best_model_{args.task}")
best_reward = -1000000.0
model_loaded_ok = False
# Architecture mismatch check
if os.path.exists(load_path + "_actor.pth"):
try:
print(f"Attempting to resume from: {load_path}")
# Note: SAC architecture is different from TDD-ND.
# This will likely fail if loading an old TDD-ND checkpoint.
agent.load(load_path)
print("Successfully loaded model weights.")
model_loaded_ok = True
except Exception as e:
print(f"Architecture Mismatch: {e}")
print("Old TDD-ND weights incompatible with SAC. Starting fresh.")
model_loaded_ok = False
# Load metadata only if model loaded successfully (avoids stale TDD-ND rewards)
task_meta_path = os.path.join(_MODELS_DIR, f"best_model_{args.task}_meta.json")
if model_loaded_ok and os.path.exists(task_meta_path):
with open(task_meta_path, "r") as f:
meta = json.load(f)
best_reward = meta.get("best_reward", -1000000.0)
agent.total_it = meta.get("total_it", 0)
print(f"Resumed best reward for task '{args.task}': {best_reward:.2f}")
replay_buffer = ReplayBuffer(config.state_dim, config.action_dim, config.buffer_size)
print(f"Starting SAC Agent Training on task: {args.task}")
os.makedirs(_MODELS_DIR, exist_ok=True)
total_steps = 0
for ep in range(args.episodes):
state = client.reset(seed=ep + int(time.time())%1000, task=args.task)
ep_reward = 0
ep_steps = 0
done = False
while not done:
# SAC uses stochastic sampling during training
if total_steps < config.exploration_steps and not os.path.exists(load_path + "_actor.pth"):
action = np.random.uniform(-config.max_action, config.max_action, size=config.action_dim)
else:
action = agent.select_action(np.array(state), evaluate=False)
next_state, reward, terminated, truncated, _ = client.step(action)
done = terminated or truncated
replay_buffer.add(state, action, next_state, reward, done)
state = next_state
ep_reward += reward
ep_steps += 1
total_steps += 1
if replay_buffer.size > config.batch_size:
agent.train(replay_buffer)
# Print metrics
alpha_val = agent.alpha.item() if hasattr(agent, 'alpha') else 0.0
print(f"Ep {ep+1}/{args.episodes} | Reward: {ep_reward:.2f} | Steps: {ep_steps} | Alpha: {alpha_val:.4f}")
# Save best model
if ep_reward > best_reward:
best_reward = ep_reward
save_path = os.path.join(_MODELS_DIR, f"best_model_{args.task}")
agent.save(save_path)
with open(os.path.join(_MODELS_DIR, f"best_model_{args.task}_meta.json"), "w") as f:
json.dump({"best_reward": best_reward, "total_it": agent.total_it, "algorithm": "SAC"}, f)
print(f" --> New best model saved! Reward: {best_reward:.2f}")
except Exception as outer_e:
print(f"Training Error: {outer_e}")
finally:
if server_process:
print("Shutting down local server...")
server_process.terminate()
else:
print("Training session finished.")
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--task", type=str, default="hard", choices=["easy", "medium", "hard"])
parser.add_argument("--episodes", type=int, default=50) # Increased default for SAC
parser.add_argument("--load-model", type=str, default=None)
args = parser.parse_args()
train(args)
|