Prit-Sudo/Organizational-Simulation-Environment

Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

AAOS: Autonomous AI Organization Simulator

OpenEnv Hackathon Submission

• —🚀 Hugging Face Space: [Insert Space URL here]
• —📝 Blog Post / Writeup: https://huggingface.co/spaces/Prit-Sudo/Organizational-Simulation-Environment/blob/main/BLOG.md
• —📝 Training Logs: https://huggingface.co/spaces/Prit-Sudo/Organizational-Simulation-Environment/blob/main/training_logs.md

1. The Problem: Why It Matters

The Capability Gap: While LLMs excel at conversational tasks, they struggle as autonomous agents in complex, multi-stakeholder organizational environments. Real-world business operations require agents to not only execute tasks but to prioritize resources, coordinate with other specialized roles, and manage compounding risks (like failing infrastructure or policy violations).

The Target Domain: Organizational decision-making is a high-impact engineering task that is severely underexplored in RL/LLM training. Existing benchmarks often rely on static QA or simple games (grid-worlds, chess). AAOS exists to teach LLMs how to navigate the messy reality of a modern tech organization—balancing immediate sprint goals against long-term system stability and compliance.

2. The Environment: What the Agent Experiences

AAOS is a deterministic, partially observable multi-agent simulation. The model plays the role of 6 specialized agents (Product Manager, Engineer, Data Analyst, Risk/Compliance, Operations, Oversight).

• —What the Agent Sees: A structured JSON observation representing a slice of the global state (e.g., the Engineer sees error rates and compute resources, while Compliance sees risk registries and policies).
• —What the Agent Does: Submits structured JSON actions via a typed API (e.g., delegate_task, approve_resource, raise_risk_flag, resolve_incident, update_policy).
• —What the Agent is Rewarded For: A rigorous, hard-to-game reward signal clamped between 0.0 and 1.0. The reward is a weighted composite of:
• —Business Success (25%): Task completion and deadlines.
• —Collaboration (20%): Cross-agent coordination.
• —Risk Management (20%): Incident resolution and policy compliance.
• —Efficiency (15%): Resource and cost management.
• —Adaptability (20%): Recovering from injected scenario shocks.

Motivation

Real organizational decision-making is a high-impact engineering task, but it is hard to evaluate with unstructured prompts alone. This project provides:

• —deterministic scenarios
• —typed API contracts
• —explicit tasks
• —deterministic graders
• —normalized scores in the range 0.0 to 1.0

That allows fair comparison between different models or agent policies.

3. Final Project Structure

aaos-env/
|-- .env
|-- .dockerignore
|-- .gitignore
|-- Dockerfile
|-- README.md
|-- client.py
|-- inference.py
|-- models.py
|-- openenv.yaml
|-- requirements.txt
|-- test_smoke.py
`-- server/
    |-- __init__.py
    |-- app.py
    |-- environment.py
    |-- tasks.py
    |-- scenarios.py
    |-- oversight.py
    `-- graders/
        |-- __init__.py
        |-- utils.py
        |-- business_grader.py
        |-- collaboration_grader.py
        `-- risk_grader.py

Core Components

server/app.py

FastAPI entrypoint. It exposes:

• —GET /health
• —POST /reset
• —POST /step
• —GET /state

server/environment.py

This is the main environment engine. It:

• —picks a scenario
• —resets state
• —executes multi-agent actions
• —updates reward
• —decides when a task is complete

server/tasks.py

Defines the available tasks:

• —org_simulation
• —crisis_management
• —full_episode

server/scenarios.py

Contains deterministic organizational scenario datasets:

• —product_launch_failure
• —resource_conflict
• —silent_data_drift
• —policy_change_shock

Each scenario includes:

• —tasks with deadlines and effort requirements
• —resource pools (compute, budget, personnel)
• —risks with severity and visibility rules
• —system metrics (latency, error_rate, uptime, cost)
• —deterministic events injected at specific timesteps
• —scenario flags tracking progression

server/graders/

Contains the deterministic scoring functions:

• —business grader (task completion + efficiency)
• —collaboration grader (coordination + adaptability)
• —risk grader (risk management + safety)

server/oversight.py

Safety auditing and penalty system. Audits every action for:

• —approve_resource without prior risk check
• —ignoring critical risks for 3+ steps
• —policy violations (resource threshold)
• —repeat offender escalation (3+ violations)

models.py

Defines typed Pydantic models for:

• —actions (AgentAction)
• —observations (Observation)
• —rewards
• —step/reset requests and responses
• —full environment state

inference.py

This is the baseline inference script used by evaluation. It:

• —calls /reset
• —reads the observation
• —chooses actions for all 6 agents
• —calls /step
• —prints structured logs in [START], [STEP], and [END] format

Tasks

1. org_simulation

Goal:

Run basic multi-agent organizational collaboration — propose plans, delegate tasks, coordinate across agents.

Difficulty:

easy

Completion:

Finishes when all tasks are delegated or in-progress/completed.

2. crisis_management

Goal:

Handle crisis scenarios — identify risks, resolve incidents, and coordinate crisis response across agents.

Difficulty:

medium

Completion:

Finishes when all detected risks are mitigated and all incidents are resolved.

3. full_episode

Goal:

Complete a full episode — collaborate, manage risks, resolve incidents, apply remediation, and adapt to dynamic events.

Difficulty:

hard

Completion:

Finishes when max_steps reached (always runs the full episode).

Observation Space

Each API step returns an Observation object with:

• —task
• —step_count
• —max_steps
• —scenario_id
• —scenario_title
• —agent_observations (per-agent partial observations with role-based visibility)
• —system_metrics
• —resources
• —active_risks
• —active_incidents
• —scenario_events
• —partial_scores
• —hints
• —degraded

This observation is the actual input the agent reasons over.

Partial Observability

Each of the 6 agents sees a DIFFERENT slice of the global state:

Action Space

The environment supports thirteen actions.

propose_plan Payload:

{"plan": "Sprint plan for step 0: prioritize critical tasks"}

accept_plan Payload:

{"plan_id": "plan_0_product_manager"}

reject_plan Payload:

{"plan_id": "plan_0_product_manager"}

delegate_task Payload:

{"task_id": "task_frontend", "to": "engineer"}

request_resource Payload:

{"amount": 15, "resource_type": "compute"}

approve_resource Payload:

{"amount": 15, "resource_type": "compute"}

query_metrics Payload:

{}

raise_risk_flag Payload:

{"risk_id": "risk_scaling", "description": "scaling_bottleneck"}

report_anomaly Payload:

{"description": "Data quality dropped to 0.65", "severity": "high"}

resolve_incident Payload:

{"incident_id": "incident_step8"}

audit_agent Payload:

{"target_agent": "engineer"}

update_policy Payload:

{"policy": "Safety review required at step 7"}

no_op Payload:

{}

Reward Logic

The reward is deterministic and clamped to 0.0 to 1.0.

Formula:

total =
  0.25 * business_success  +
  0.20 * collaboration     +
  0.20 * risk_management   +
  0.15 * efficiency        +
  0.20 * adaptability

Reward parts:

• —business_success: task completion quality and deadline adherence
• —collaboration: inter-agent coordination and communication
• —risk_management: risk detection, incident resolution, policy compliance
• —efficiency: resource utilization, cost management, latency performance
• —adaptability: response to scenario events, recovery, policy adaptation

Reward delivery:

• —60% immediate (returned on the step)
• —40% delayed (released every 5 steps + episode end)
• —Oversight penalties push reward toward 0 (never negative)

OpenEnv Metadata

The root openenv.yaml contains:

• —environment name and version
• —app entrypoint
• —HTTP runtime type
• —task metadata with difficulty

Current task metadata:

• —org_simulation
• —crisis_management
• —full_episode

Actual Input to the System

The system input is not a single plain English prompt. The real input is the structured observation returned by /reset or /step.

Example observation:

{
  "task": "full_episode",
  "step_count": 0,
  "max_steps": 50,
  "scenario_id": "product_launch_failure",
  "scenario_title": "Aggressive deadline with hidden scaling issues",
  "agent_observations": {
    "product_manager": {
      "role": "product_manager",
      "step": 0,
      "visible_tasks": [
        {"id": "task_frontend", "name": "frontend_polish", "status": "pending", "priority": "high", "deadline_step": 10, "quality": 0.0, "effort_remaining": 4}
      ],
      "system_metrics": {"latency": 200, "uptime": 0.999},
      "resources": "hidden",
      "risks": "hidden"
    },
    "engineer": {
      "role": "engineer",
      "step": 0,
      "visible_tasks": [
        {"id": "task_backend", "name": "backend_scaling", "status": "pending", "priority": "critical"}
      ],
      "system_metrics": {"latency": 200, "error_rate": 0.02},
      "resources": {"compute": 80},
      "risks": []
    }
  },
  "system_metrics": {"latency": 200, "error_rate": 0.02, "uptime": 0.999, "cost": 0.0},
  "resources": {"compute": 80, "budget": 50, "personnel": 5},
  "active_risks": [],
  "active_incidents": [],
  "scenario_events": [],
  "partial_scores": {
    "business": 0.0,
    "collaboration": 0.0,
    "risk": 0.0,
    "efficiency": 1.0,
    "adaptability": 0.0
  },
  "hints": ["Complete all objectives: collaborate, manage risks, resolve incidents, and adapt to events."],
  "degraded": false
}

Actual Output from the System

The system output is a sequence of actions and scores.

Example actions (multi-agent, one per agent per step):

{
  "product_manager": {"action_type": "propose_plan", "payload": {"plan": "Sprint plan for step 0: prioritize critical tasks"}},
  "engineer": {"action_type": "delegate_task", "payload": {"task_id": "task_backend", "to": "operations"}},
  "data_analyst": {"action_type": "query_metrics", "payload": {}},
  "risk_compliance": {"action_type": "query_metrics", "payload": {}},
  "operations": {"action_type": "approve_resource", "payload": {"amount": 15, "resource_type": "compute"}},
  "oversight": {"action_type": "audit_agent", "payload": {"target_agent": "engineer"}}
}

Final terminal output format:

[START] task=full_episode env=aaos-env model=Qwen/Qwen2.5-72B-Instruct backend=https://router.huggingface.co/v1 seed=42
[STEP] step=1 action={"action_type": "propose_plan", "payload": {"plan": "Sprint plan..."}} reward=0.3210 done=False error=None
...
[END] success=True steps=25 score=0.7281 rewards=[...]

This format is important because the evaluation pipeline expects strict [START], [STEP], and [END] lines.

Baseline Behavior

The baseline inference script runs all 3 tasks sequentially:

1. 1.org_simulation
2. 2.crisis_management
3. 3.full_episode

4. Results: Untrained Baseline vs. Trained Agent

We evaluated the performance of an untrained heuristic/zero-shot baseline against our agent trained using a two-stage pipeline: Supervised Fine-Tuning (SFT) via Unsloth, followed by Group Relative Policy Optimization (GRPO) using trl.

Quantitative Comparison

Training Evidence

The plots below demonstrate the end-to-end learning process connecting the environment to the training loop.

Loss Curves (SFT + GRPO)

[image] Figure 1: SFT loss drops significantly in the first epoch. GRPO RL training maintains a stable, low loss (~0.01-0.04) while exploring policy improvements.

Reward Curves (GRPO Phase)

[image] Figure 2: Total reward (top) climbs from an initial 1.73 to converge at ~4.00 (out of 5.0 max). The per-component breakdown (bottom) shows the agent rapidly mastering formatting, action validity, and safety, while environment-specific strategic rewards (Env Score) represent the hardest capability gap.

Qualitative Improvements

After training, the agent stopped proposing empty or cyclical plans. Instead of merely reacting to incidents, the trained agent proactively queried metrics, raised risk flags before systems failed, and correctly routed compute approvals to the Operations role.

5. Training Pipeline

The project includes an end-to-end training pipeline combining Supervised Fine-Tuning (SFT) and Reinforcement Learning via GRPO.

1. Dataset Collection

Run the environment with a strong teacher model to collect raw execution episodes:

python collect_training_data.py --episodes 4 --steps 10 --model-id Qwen/Qwen2.5-7B-Instruct --token <YOUR_HF_TOKEN>

Outputs to: `training_data/raw_episodes.jsonl`

2. SFT Dataset Preparation

Format and filter the raw episodes, keeping only high-reward actions to train the SFT baseline:

python prepare_sft_dataset.py

Outputs to: `training_data/sft_dataset.jsonl`

3. Supervised Fine-Tuning (SFT)

Train the initial policy using Unsloth for fast, memory-efficient LoRA tuning:

python -m train_unsloth_sft \
   --dataset training_data/sft_dataset.jsonl \
   --output fine_tuned_model \
   --epochs 1 \
   --batch-size 1 \
   --grad-accum 8 \
   --max-seq-len 1024 \
   --lora-rank 8

Outputs to: `fine_tuned_model/`

4. Reinforcement Learning (GRPO)

Optimize the SFT model further using Group Relative Policy Optimization (GRPO) to align with environment rewards:

python -m train_grpo_rl \
    --base-model fine_tuned_model \
    --epochs 1 \
    --batch-size 1 \
    --num-generations 2 \
    --max-prompts 150

Outputs to: `aaos_grpo_model/`

You can find more detailed logs and intermediate steps from the reference training run in post_training.md.

6. Setup and Run

From the repository root:

python -m venv .venv
.venv\Scripts\Activate.ps1
pip install -r requirements.txt

Start the server:

uvicorn server.app:app --host 0.0.0.0 --port 7860

Run inference:

python inference.py

Run tests:

python -m unittest test_smoke -v

7. Required Environment Variables

The project supports:

• —API_BASE_URL
• —MODEL_NAME
• —HF_TOKEN
• —ENV_BASE_URL
• —RUN_SEED
• —MAX_STEPS
• —SUCCESS_THRESHOLD

Validation Status

The project is structured for OpenEnv-style evaluation and includes:

• —typed models
• —root inference.py
• —root openenv.yaml
• —root Dockerfile
• —3 tasks
• —graders
• —deterministic scoring
• —structured logs

This makes the project ready for final validation and submission.