| """ |
| ARYA-X GRPO Training Script |
| Run on HF Spaces A10G: ~15 min for 200 steps |
| Saves LoRA adapter to HF Hub repo specified by HF_REPO env var. |
| """ |
| import os, json, random, torch |
| from collections import defaultdict |
|
|
| |
| from unsloth import FastLanguageModel |
|
|
| MODEL_NAME = "unsloth/Llama-3.2-3B-Instruct-bnb-4bit" |
| print(f"Loading {MODEL_NAME}...") |
|
|
| model, tokenizer = FastLanguageModel.from_pretrained( |
| model_name=MODEL_NAME, |
| max_seq_length=512, |
| load_in_4bit=True, |
| dtype=None, |
| ) |
| model = FastLanguageModel.get_peft_model( |
| model, |
| r=4, |
| lora_alpha=8, |
| target_modules=["q_proj", "v_proj"], |
| lora_dropout=0.0, |
| bias="none", |
| use_gradient_checkpointing="unsloth", |
| ) |
| print("✅ Model ready.") |
|
|
| |
| SENSOR_TYPES = ["satellite", "drone", "radar"] |
| TARGET_TYPES = ["strategic", "kinetic", "airspace"] |
| CAPABILITY_MATRIX = { |
| ("satellite", "strategic"): 0.95, ("satellite", "kinetic"): 0.40, |
| ("satellite", "airspace"): 0.60, ("drone", "kinetic"): 0.95, |
| ("drone", "strategic"): 0.30, ("drone", "airspace"): 0.50, |
| ("radar", "airspace"): 0.95, ("radar", "kinetic"): 0.65, |
| ("radar", "strategic"): 0.45, |
| } |
|
|
| def make_env(seed=0): |
| rng = random.Random(seed) |
| sensors = [ |
| {"id": f"S{i+1}", "type": rng.choice(SENSOR_TYPES), |
| "available": True, "range": rng.randint(100, 500)} |
| for i in range(rng.randint(3, 5)) |
| ] |
| return {"sensors": sensors, "seed": seed, "_rng": rng} |
|
|
| def env_reset(env_state): |
| rng = env_state["_rng"] |
| targets = [ |
| {"id": f"T0_{i+1}", "priority": rng.randint(1, 3), |
| "type": rng.choice(TARGET_TYPES), "active": True} |
| for i in range(rng.randint(2, 4)) |
| ] |
| for s in env_state["sensors"]: |
| s["available"] = True |
| env_state.update({"targets": targets, "current_step": 0, "total_reward": 0.0}) |
| return env_state |
|
|
| |
| import datasets as hf_datasets |
| from trl import GRPOTrainer, GRPOConfig |
|
|
| def make_prompt(seed): |
| rng = random.Random(seed) |
| sensors = [ |
| {"id": f"S{i+1}", "type": rng.choice(SENSOR_TYPES), "range": rng.randint(100, 500)} |
| for i in range(rng.randint(3, 5)) |
| ] |
| targets = [ |
| {"id": f"T0_{i+1}", "priority": rng.randint(1, 3), "type": rng.choice(TARGET_TYPES)} |
| for i in range(rng.randint(2, 4)) |
| ] |
| sensor_lines = "\n".join(f" {s['id']} type={s['type']} range={s['range']}km" for s in sensors) |
| target_lines = "\n".join(f" {t['id']} priority={t['priority']} type={t['type']}" for t in targets) |
| sensor_ids = [s["id"] for s in sensors] |
| target_ids = [t["id"] for t in targets] |
| prompt = ( |
| f"You are an ISR sensor allocation agent.\n" |
| f"Assign ONE sensor to ONE target. Priority 3=HIGH, 2=MED, 1=LOW.\n" |
| f"Cover the highest priority target first.\n\n" |
| f"Available sensors:\n{sensor_lines}\n\n" |
| f"Active targets:\n{target_lines}\n\n" |
| f"Respond ONLY with valid JSON using these exact IDs:\n" |
| f"sensors={sensor_ids} targets={target_ids}\n\n" |
| f"{{\"sensor_id\": \"S1\", \"target_id\": \"T0_1\"}}" |
| ) |
| return prompt, sensor_ids, target_ids, targets |
|
|
| prompt_list = [] |
| metadata_list = [] |
| for seed in range(500): |
| prompt, sensor_ids, target_ids, targets = make_prompt(seed) |
| prompt_list.append(prompt) |
| metadata_list.append({"sensor_ids": sensor_ids, "target_ids": target_ids, |
| "targets": targets, "seed": seed}) |
|
|
| train_dataset = hf_datasets.Dataset.from_dict({"prompt": prompt_list}) |
| print(f"✅ Dataset ready: {len(train_dataset)} prompts") |
|
|
| |
| NUM_GENERATIONS = 8 |
|
|
| def arya_reward_func(completions, prompts=None, **kwargs): |
| rewards = [] |
| for i, completion in enumerate(completions): |
| try: |
| prompt_idx = i // NUM_GENERATIONS |
| meta = metadata_list[prompt_idx % len(metadata_list)] |
|
|
| text = completion if isinstance(completion, str) else str(completion) |
| start = text.find("{") |
| end = text.rfind("}") + 1 |
| if start == -1 or end == 0: |
| rewards.append(-2.0); continue |
|
|
| data = json.loads(text[start:end]) |
| sensor_id = data.get("sensor_id", "") |
| target_id = data.get("target_id", "") |
|
|
| if sensor_id not in meta["sensor_ids"] or target_id not in meta["target_ids"]: |
| rewards.append(-1.0); continue |
|
|
| target = next((t for t in meta["targets"] if t["id"] == target_id), None) |
| env_s = env_reset(make_env(seed=meta["seed"])) |
| sensor = next((s for s in env_s["sensors"] if s["id"] == sensor_id), None) |
| if not target or not sensor: |
| rewards.append(-1.0); continue |
|
|
| cap = CAPABILITY_MATRIX.get((sensor["type"], target.get("type", "strategic")), 0.0) |
| p = target["priority"] |
| reward = (3.0 if cap >= 0.85 else 2.0) if p == 3 else {2: 1.0, 1: 0.5}.get(p, 0.0) |
| if cap >= 0.85: |
| reward += 0.5 |
| rewards.append(float(reward)) |
| except Exception: |
| rewards.append(-2.0) |
| return rewards |
|
|
| |
| _test = arya_reward_func(['{"sensor_id": "S1", "target_id": "T0_1"}'] * 8) |
| print(f"✅ Reward sanity check: {_test}") |
|
|
| |
| grpo_args = GRPOConfig( |
| output_dir="./grpo_out", |
| num_train_epochs=1, |
| per_device_train_batch_size=8, |
| gradient_accumulation_steps=2, |
| max_steps=200, |
| learning_rate=2e-5, |
| warmup_steps=10, |
| logging_steps=10, |
| save_steps=50, |
| report_to="none", |
| ) |
|
|
| grpo_trainer = GRPOTrainer( |
| model=model, |
| args=grpo_args, |
| tokenizer=tokenizer, |
| reward_funcs=arya_reward_func, |
| train_dataset=train_dataset, |
| ) |
|
|
| print(f"\nStarting training — 200 steps on A10G (~15 min)...\n") |
| train_result = grpo_trainer.train() |
| print(f"\n✅ Training complete. Loss: {train_result.training_loss:.4f}") |
|
|
| |
| FastLanguageModel.for_inference(model) |
|
|
| total_reward, conflicts, valid_outputs = 0.0, 0, 0 |
| for seed in range(50): |
| prompt, sensor_ids, target_ids, targets = make_prompt(seed) |
| inputs = tokenizer(prompt, return_tensors="pt").to("cuda") |
| with torch.no_grad(): |
| outputs = model.generate(**inputs, max_new_tokens=64, do_sample=False, |
| pad_token_id=tokenizer.eos_token_id) |
| new_ids = outputs[0][inputs["input_ids"].shape[-1]:] |
| text = tokenizer.decode(new_ids, skip_special_tokens=True) |
| start, end = text.find("{"), text.rfind("}") + 1 |
| try: |
| data = json.loads(text[start:end]) |
| sensor_id = data.get("sensor_id", "") |
| target_id = data.get("target_id", "") |
| if sensor_id not in sensor_ids or target_id not in target_ids: |
| conflicts += 1; total_reward -= 1.0; continue |
| target = next((t for t in targets if t["id"] == target_id), None) |
| env_s = env_reset(make_env(seed=seed)) |
| sensor = next((s for s in env_s["sensors"] if s["id"] == sensor_id), None) |
| if target and sensor: |
| cap = CAPABILITY_MATRIX.get((sensor["type"], target.get("type", "strategic")), 0.0) |
| p = target["priority"] |
| reward = (3.0 if cap >= 0.85 else 2.0) if p == 3 else {2: 1.0, 1: 0.5}.get(p, 0.0) |
| total_reward += reward; valid_outputs += 1 |
| else: |
| conflicts += 1 |
| except Exception: |
| conflicts += 1; total_reward -= 2.0 |
|
|
| results = { |
| "avg_reward": round(total_reward / 50, 3), |
| "conflict_rate": round(conflicts / 50, 3), |
| "coordination_score": round(1.0 - conflicts / 50, 3), |
| "valid_outputs": valid_outputs, |
| } |
| print(f"\n{'='*45}") |
| print(f" TRAINED MODEL RESULTS") |
| print(f"{'='*45}") |
| for k, v in results.items(): |
| print(f" {k:<22}: {v}") |
| print(f"{'='*45}") |
|
|
| |
| os.makedirs("checkpoints/arya_x_lora", exist_ok=True) |
| model.save_pretrained("checkpoints/arya_x_lora") |
| tokenizer.save_pretrained("checkpoints/arya_x_lora") |
|
|
| metrics = {**results, "training_loss": train_result.training_loss, |
| "steps": 200, "model": "Llama-3.2-3B-Instruct", "method": "GRPO + LoRA r=4"} |
| os.makedirs("logs", exist_ok=True) |
| with open("logs/training_metrics.json", "w") as f: |
| json.dump(metrics, f, indent=2) |
| print("✅ Saved: logs/training_metrics.json") |
|
|
| HF_REPO = os.environ.get("HF_REPO", "") |
| HF_TOKEN = os.environ.get("HF_TOKEN", "") |
| if HF_REPO and HF_TOKEN: |
| from huggingface_hub import login |
| login(token=HF_TOKEN) |
| model.push_to_hub(HF_REPO, token=HF_TOKEN) |
| tokenizer.push_to_hub(HF_REPO, token=HF_TOKEN) |
| print(f"✅ Pushed to Hub: {HF_REPO}") |
| else: |
| print("ℹ️ Set HF_REPO and HF_TOKEN env vars to push to Hub.") |
|
|
