File size: 6,636 Bytes
bda28c2 | 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 146 147 | """
AegisGym CPU-Optimized RL Training (REINFORCE)
Manual policy gradient loop for training without a GPU.
"""
import os
import json
import torch
import torch.nn.functional as F
from transformers import AutoTokenizer, AutoModelForCausalLM
from client_env import get_sync_client
from server import AegisGymEnv
from models import AuditAction
from datasets import load_dataset
import itertools
# ─── Config ──────────────────────────────────────────────────────────────────
MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-0.5B-Instruct")
ENV_URL = os.getenv("AEGISGYM_URL", "https://armaan020-aegisgym.hf.space")
LR = 1e-5
EPISODES = 5
ACCUMULATION_STEPS = 4
SYSTEM_PROMPT = """You are a financial compliance auditor AI.
Given a transaction scenario and regulatory context, respond with a JSON object:
{"action_type": "FLAG|APPROVE|REQUEST_INFO", "target_id": "<account_id>", "regulation_citation": "<regulation>"}
Be precise and concise."""
# ─── Dataset ──────────────────────────────────────────────────────────────────
print(f"Loading dataset: SecureFinAI-Lab/Regulations_QA...")
try:
raw_dataset = load_dataset("SecureFinAI-Lab/Regulations_QA", split="train", streaming=True)
dataset_iterator = itertools.cycle(iter(raw_dataset))
print("Dataset loaded via streaming.")
except Exception as e:
print(f"Dataset loading failed: {e}. Using generic prompts.")
dataset_iterator = None
# ─── Model & Opt ─────────────────────────────────────────────────────────────
print(f"Loading {MODEL_NAME} on CPU...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype="auto", device_map="cpu")
optimizer = torch.optim.AdamW(model.parameters(), lr=LR)
model.train()
# ─── Action Parser ────────────────────────────────────────────────────────────
def parse_action(text: str) -> AuditAction:
try:
start = text.find("{")
end = text.rfind("}") + 1
if start >= 0 and end > start:
data = json.loads(text[start:end])
return AuditAction(**data)
except Exception:
pass
return AuditAction(action_type="REQUEST_INFO", target_id="UNKNOWN", regulation_citation="parse_error")
# ─── Training Loop ────────────────────────────────────────────────────────────
def train():
print(f"\n=== Starting CPU training (REINFORCE) ===")
env = get_sync_client(ENV_URL)
total_reward = 0
for i in range(EPISODES):
print(f"Episode {i+1}/{EPISODES}...")
# Reset Env
obs_payload = env.reset()
obs = obs_payload.get("observation", {})
# Build Prompt (with dataset augmentation)
dataset_context = ""
custom_prompt = "Audit the following transaction."
if dataset_iterator:
item = next(dataset_iterator)
custom_prompt = item.get("question", custom_prompt)
dataset_context = f"\nRegulatory Context: {item.get('answer', '')}"
tier = env.state().get("current_tier", "easy")
user_msg = (
f"{custom_prompt}\n\n"
f"Tier: {tier.upper()}\n"
f"Transactions: {obs.get('transactions', [])}\n"
f"Context: {obs.get('retrieved_regs', '')} {dataset_context}\n"
f"Account: {obs.get('account_metadata', {})}"
)
messages = [
{"role": "system", "content": SYSTEM_PROMPT},
{"role": "user", "content": user_msg},
]
prompt = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
# Tokenize and Generate with Gradients
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
# We need to compute log-probs of the generated completion
# and use those for the REINFORCE loss.
# This is high-memory for long sequences, so we keep max_new_tokens small.
outputs = model.generate(**inputs, max_new_tokens=128, return_dict_in_generate=True, output_scores=True)
# Re-run full forward pass on the generated completion to get log_probs with gradients
# The completion starts after inputs.input_ids.shape[1]
full_ids = outputs.sequences
context_len = inputs.input_ids.shape[1]
completion_ids = full_ids[:, context_len:]
logits = model(full_ids).logits
# Align logits with completion ids (logits[i] is logit for token full_ids[i+1])
completion_logits = logits[:, context_len-1:-1, :]
log_probs = F.log_softmax(completion_logits, dim=-1)
# Gather the log-probs of the actual tokens chosen
selected_log_probs = log_probs.gather(2, completion_ids.unsqueeze(-1)).squeeze(-1)
# Sum log-probs for the sequence (per REINFORCE)
episode_log_prob = selected_log_probs.sum()
# Execution
completion_text = tokenizer.decode(completion_ids[0], skip_special_tokens=True)
action = parse_action(completion_text)
# Environment Step
result = env.step(action.model_dump())
reward = float(result.get("reward", 0.0))
print(f" Actions: {action.action_type} | Reward: {reward}")
total_reward += reward
# Loss: -Sum(log_probs) * Reward
# Minimize this to maximize (log_probs * Reward)
loss = -episode_log_prob * reward
loss = loss / ACCUMULATION_STEPS # gradient accumulation
loss.backward()
if (i+1) % ACCUMULATION_STEPS == 0:
optimizer.step()
optimizer.zero_grad()
print(" --- Weights Updated! ---")
print(f"\nTraining Complete. Avg Reward: {total_reward/EPISODES}")
model.save_pretrained("aegisgym-cpu-agent")
tokenizer.save_pretrained("aegisgym-cpu-agent")
print(f"Model saved to 'aegisgym-cpu-agent'")
if __name__ == "__main__":
train()
|