Apppublic

22f3001825/AegisWeave-Training-A100

sourceHugging Faceupdated 5mo agoView on Hugging Face

0likes

App README

AegisWeave 🛡️

Meta PyTorch · OpenEnv · Hackathon India 2026

Team: Mayank & Ali · Model: Qwen2.5-3B + GRPO + Unsloth

AegisWeave is an OpenEnv-compliant RL benchmark where an AI assistant manages a 72-hour executive product launch crisis across 50 decision steps. It coordinates 5 stakeholders with hidden emotional states, adapts to stochastic API schema drift detected only through indirect signals, and operates under oversight from a fine-tuned 0.5B Monitor LLM.

Hackathon Themes

Theme	Track
Theme 2	Long-Horizon Planning — 50-step episodes, cascade chains requiring 20+ step lookahead
Theme 3.2	Personalized Tasks — Executive assistant with personal/work/family conflicts
Patronus AI	Schema Drift — Poisson 8%/step, zero notification, online adaptation
Fleet AI	Scalable Oversight — SFT fine-tuned 0.5B Monitor with ALLOW/VETO decisions

🛠️ Recent Fixes for 100% Correct Training

Before executing the A100 training, we thoroughly audited and fixed the environment and training pipeline to ensure a 100% mathematically correct RL setup:

—GRPO State-Replay (Critical): Fixed a bug where TRL's GRPOTrainer evaluated rewards from step-0 regardless of the prompt depth. Prompts now inject seed and past_actions into the dataset, and the reward function perfectly replays the exact action history.
—Reward Shaping: Added a small, dense shaping reward (+0.1) for targeting the stakeholder with the lowest satisfaction, solving stalled learning in early Stage 1.
—Agent Logic: Corrected suggest_break logic paradoxes and implemented dynamic pr_approved checks to gracefully handle simulated Patronus AI schema drift.
—Curriculum Robustness: Expanded the prompt dataset size to 256 for greater diversity and fixed brittle curriculum advancement loops.

🚀 Quick Start

1. Installation

bash

pip install -r requirements.txt

2. Local Validation (CPU — No Docker)

Run a dry-run validation of the environment, agents, and curriculum:

bash

python -m training.grpo_train --dry-run

This validates the full pipeline (environment, agents, reward function, curriculum stages, trajectory generator) without GPU.

3. Run with Docker (OpenEnv Server)

bash

docker build -t aegisweave .
docker run -p 8000:8000 aegisweave

Open http://localhost:8000/web for the OpenEnv Web UI.

📊 Training Pipeline (A100 Required)

Full pipeline uses Unsloth + TRL GRPO on A100 40GB (~$60 total, well within $200 budget).

Step 1: Generate Expert Trajectories

Run rule-based expert policy to generate ~10,000 labeled (state, action, harmful) pairs for Monitor SFT:

bash

python -m training.generate_trajectories --n-episodes 200

Step 2: Fine-tune Monitor Agent (Fleet AI)

Train the 0.5B Monitor on labeled trajectories (~30 mins, ~$3):

bash

python -m training.monitor_sft

Step 3: GRPO Curriculum Training

3-stage curriculum producing 3 visible reward uptrends:

Stage	Stakeholders	Drift	Monitor	Steps
1 — Basic Planning	2	❌	❌	300
2 — Drift Adaptation	4	✅	❌	300
3 — Full Environment	5	✅	✅	400

bash

python -m training.grpo_train --stage 1

Step 4: Evaluate

Compare zero-shot baseline vs. trained agent:

bash

python -m training.evaluate --n-episodes 50 --trained-checkpoint outputs/stage3_checkpoint

🏗 Architecture

AegisWeave/
├── environment/
│   ├── env.py              # Main OpenEnv + Gymnasium environment
│   ├── stakeholders.py     # 5 NPCs with hidden state + curriculum filtering
│   ├── cascades.py         # 5 consequence chains (C1-C5)
│   ├── schema_drift.py     # Poisson drift engine (3 corruption types)
│   ├── reward.py           # Dense step + sparse terminal rewards
│   ├── config.py           # All tunable constants centralized
│   ├── models.py           # Pydantic observation/action models
│   ├── state.py            # ExecutiveState, ObservationSpace, HiddenState
│   ├── actions.py          # 9 action types + validation
│   ├── client.py           # OpenEnv client connector
│   └── server/             # FastAPI server for Docker deployment
├── agents/
│   ├── assistant_agent.py  # Qwen2.5-3B + rule-based fallback
│   └── monitor_agent.py    # Qwen2.5-0.5B overseer (SFT fine-tuned)
├── training/
│   ├── grpo_train.py       # Main GRPO curriculum training script
│   ├── curriculum.py       # 3-stage difficulty manager
│   ├── generate_trajectories.py  # Expert trajectory generation
│   ├── monitor_sft.py      # Monitor SFT pipeline
│   └── evaluate.py         # Before/after metrics + reward curves
├── demo/
│   └── episode_visualizer.py  # HTML + ASCII episode timeline for pitch
├── Dockerfile              # OpenEnv container spec
├── requirements.txt
└── colab_training.ipynb    # One-click training notebook

🌐 Deployment

bash

openenv push

👥 Team

—Mayank — Environment core, OpenEnv compliance, Docker, HuggingFace Spaces
—Ali — Monitor agent, training pipeline, GRPO curriculum, Colab notebook