AdamK29/Meta-OpenENV-Hackathon

1import torch2import torch.nn as nn3import torch.optim as optim4import asyncio5 6from client import EmailEnvClient7from core.models import EmailAction8 9import ssl10import certifi11 12ssl._create_default_https_context = ssl._create_unverified_context13 14# ---------- ACTION SPACE ----------15 16ACTIONS = ["spam", "important", "urgent"]17 18 19# ---------- POLICY NETWORK ----------20 21class PolicyNet(nn.Module):22    def __init__(self):23        super().__init__()24 25        self.fc = nn.Sequential(26            nn.Linear(10, 64),27            nn.ReLU(),28            nn.Linear(64, 32),29            nn.ReLU(),30            nn.Linear(32, len(ACTIONS)),31            nn.Softmax(dim=-1)32        )33 34    def forward(self, x):35        return self.fc(x)36 37 38# ---------- STATE ENCODER ----------39 40def encode_state(obs):41    text = obs.email_text.lower()42 43    features = [44        int("free" in text),45        int("win" in text),46        int("urgent" in text),47        int("asap" in text),48        int("deadline" in text),49        int("meeting" in text),50        int("offer" in text),51        int("client" in text),52        int("server" in text),53        int("issue" in text),54    ]55 56    return torch.tensor(features, dtype=torch.float32)57 58 59# ---------- TRAIN LOOP ----------60 61async def train():62 63    env = EmailEnvClient(base_url="https://adamk29-meta-openenv-hackathon.hf.space")64    await env.__aenter__()65 66    model = PolicyNet()67    optimizer = optim.Adam(model.parameters(), lr=0.003)68 69    EPISODES = 5070    GAMMA = 0.9971    ENTROPY_BETA = 0.0172 73    for episode in range(EPISODES):74 75        result = await env.reset(task="medium")76 77        log_probs = []78        rewards = []79        entropies = []80 81        while True:82 83            state = encode_state(result.observation)84 85            probs = model(state)86 87            # Add exploration noise88            probs = probs + 1e-689            probs = probs / probs.sum()90 91            dist = torch.distributions.Categorical(probs)92 93            action_idx = dist.sample()94            action = ACTIONS[action_idx.item()]95 96            result = await env.step(97                EmailAction(action_type="classify", content=action)98            )99 100            reward = result.reward or 0.0101 102            log_probs.append(dist.log_prob(action_idx))103            rewards.append(reward)104            entropies.append(dist.entropy())105 106            if result.done:107                break108 109        # ---------- COMPUTE DISCOUNTED RETURNS ----------110 111        returns = []112        G = 0113 114        for r in reversed(rewards):115            G = r + GAMMA * G116            returns.insert(0, G)117 118        returns = torch.tensor(returns)119 120        # Normalize returns (VERY IMPORTANT)121        returns = (returns - returns.mean()) / (returns.std() + 1e-8)122 123        # ---------- LOSS ----------124 125        policy_loss = []126        entropy_loss = []127 128        for log_prob, R, entropy in zip(log_probs, returns, entropies):129            policy_loss.append(-log_prob * R)130            entropy_loss.append(-ENTROPY_BETA * entropy)131 132        loss = torch.stack(policy_loss).sum() + torch.stack(entropy_loss).sum()133 134        optimizer.zero_grad()135        loss.backward()136        optimizer.step()137 138        total_reward = sum(rewards)139 140        print(f"Episode {episode} | Reward: {total_reward:.2f}")141 142    # ---------- SAVE MODEL ----------143 144    torch.save(model.state_dict(), "email_agent.pth")145    print("✅ Model saved as email_agent.pth")146 147    await env.__aexit__(None, None, None)148 149 150if __name__ == "__main__":151    asyncio.run(train())