Skip to content

garg-tejas/dispatchR

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

144 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

title DispatchR Environment Server
emoji 🚑
colorFrom red
colorTo blue
sdk docker
pinned false
tags
openenv
emergency-response
resource-allocation
long-horizon

DispatchR

DispatchR — where Reinforcement Learning meets emergency Response. Training AI dispatchers to make life-or-death decisions under uncertainty, radio silence, and chaos.

A reinforcement learning environment for emergency medical dispatch, built for the Meta OpenEnv Hackathon (India, April 2026).

Status: 47/47 tests passing. Training pipeline validated on L4. Ready for HF Hub Jobs.

Links

What It Is

DispatchR simulates an 8-hour shift for an emergency dispatch commander in a 20-node city. The agent must:

  • Dispatch 6 units to cardiac, trauma, and fire emergencies
  • Deal with radio delays (statuses lag 2-3 steps behind reality)
  • Verify ghost calls (AI-generated false emergencies)
  • Respond to city events (bridge collapse, heatwave, hospital divert, unit breakdown)
  • Maintain coverage across 5 zones so no area is left uncovered

The environment is a partially-observable Markov decision process (POMDP) with hidden severity modifiers, noisy caller tones, and dynamic traffic conditions.

Hackathon Theme Alignment

DispatchR maps to three themes simultaneously, maximizing the innovation score:

Primary: Theme #2 — (Super) Long-Horizon Planning & Instruction Following

80-step episodes with sparse end-of-episode rewards. The agent must plan across the full shift, not just react to the current call. Early dispatch mistakes compound — sending the wrong unit to a low-priority call means that unit is unavailable when a cardiac emergency arrives 20 steps later. The agent must decompose the shift into sub-goals (coverage maintenance, triage, event response) and recover from early errors.

Secondary: Theme #3.1 — World Modeling (Professional Tasks)

Partial observability forces genuine world-model construction:

  • Delayed radio updates: Unit statuses lag 2-3 steps behind reality
  • Hidden severity: Caller tone (calm/agitated/screaming) is observable; true severity is not
  • Ghost calls: AI-generated false emergencies with noisy verification
  • Dynamic traffic: Bridge collapses reroute all units in real-time
  • Noisy hospital capacity: Reported status differs from true capacity

The agent cannot succeed by pattern-matching — it must maintain a persistent belief state about the true world.

Tertiary: Theme #4 — Self-Improvement (Adaptive Curriculum)

Performance-gated curriculum learning auto-escalates difficulty:

  • Warmup (3 calls, 0% false alarms) → Expert (12 calls, 20% false alarms, 40% event rate)
  • Escalation triggered only when mean batch reward stabilizes above 0.65
  • The environment itself gets harder as the agent improves, preventing plateauing on easy scenarios

Environment Features

Feature Description
20-Node City 5 zones (Downtown, Highway, Hills, Industrial, Suburbs) with Dijkstra shortest-path routing
6 Units State machine: idle → en_route → on_scene → returning
Radio Delay 10% chance status updates lag 2-3 steps (POMDP)
Caller Panic Bias Hidden panic modifier distorts reported severity; calm voice can mask cardiac
False Alarms 10-20% of calls are false alarms (infinite deadline, no penalty for ignoring)
Ghost Calls AI-generated deepfake emergencies (0-10%). verify action returns noisy confidence
City Events Bridge collapse (reroutes traffic), hospital divert, heatwave (2x cardiac rate), unit breakdown
Hospital Capacity Hidden true capacity; reported status is noisy
Adversarial Designer Weakness tracker biases call generation toward agent failure zones
Adaptive Curriculum Auto-escalates from Warmup (3 calls, 0 events) to Expert (12 calls, 40% event rate). Now controllable via CLI in training
LLM-Friendly Prompts Active Calls section explicitly shows (none) when empty; system prompt includes inline examples and conciseness rules to reduce rambling

Architecture

dead-air/
├── README.md                       # This file
├── pyproject.toml                  # Package dependencies
│
├── train_trl_grpo.py               # PRIMARY: TRL-native GRPO training (recommended)
├── train_unsloth_grpo.py           # LEGACY: Unsloth GRPO training path
├── train_grpo.py                   # LEGACY: Manual GRPO training path
├── eval.py                         # Greedy baseline evaluation
├── demo.py                         # Interactive terminal demo
├── diagnose.py                     # Per-episode diagnostic script
├── inference.py                    # Run trained checkpoint (supports Unsloth loading)
├── export_episode.py               # Export greedy episode to JSON for visualization
├── visualize_env.py                # Render episode JSON → MP4 animation
├── plot_rewards.py                 # Reward curve visualization
├── client.py                       # OpenEnv WebSocket client
├── models.py                       # Pydantic Action/Observation schemas
│
├── scripts/
│   ├── launch_hf_job.py            # HF Hub Jobs launcher (Python CLI)
│   ├── launch_hf_job.sh            # HF Hub Jobs launcher (bash)
│   ├── plot_training.py            # Plot training metrics from metrics.json
│   └── visualize_city.py           # City graph visualization tools
│
├── server/
│   ├── app.py                      # FastAPI + WebSocket OpenEnv server
│   ├── dispatcher_environment.py   # Core env: reset, step, reward
│   ├── grpo_env_wrapper.py         # Wrapper for LLM training (JSON parser + metrics)
│   ├── training_tracker.py         # Real-time progress tracking + CSV + plots
│   ├── unsloth_grpo_utils.py       # GRPO loss computation for batched token trajectories
│   ├── city_graph.py               # NetworkX graph + Dijkstra oracle
│   ├── call_generator.py           # Poisson arrivals + panic + false alarms + ghosts
│   ├── unit_model.py               # Unit state machine + RadioDelayBuffer
│   ├── traffic_model.py            # Time-varying edge weights + accident injection
│   ├── hospital_model.py           # Hidden capacity + noisy divert signals
│   ├── event_scheduler.py          # City event templates + random triggers
│   ├── adversarial_designer.py     # Weakness tracker + dynamic bias
│   ├── curriculum.py               # Auto-escalate/de-escalate difficulty
│   ├── reward.py                   # 3-component episode reward + action validity shaping
│   └── constants.py                # Medical deadlines, city topology, unit configs
│
└── tests/
    ├── test_environment.py         # Env lifecycle (reset, step, dispatch, hold)
    ├── test_call_generator.py      # Call generation, panic, false alarms, ghosts
    ├── test_unit_model.py          # Unit state machine + radio delay buffer
    ├── test_city_graph.py          # Graph topology + oracle assignment
    ├── test_hospital_events.py     # Hospital model + event scheduler
    └── test_connection.py          # End-to-end episode + curriculum

Quick Start

1. Run Tests

python -m pytest tests/
# 47 passed in ~0.5s

2. Greedy Baseline

python eval.py

Runs 10 episodes with a greedy dispatcher (closest idle unit to highest-priority call). Expected mean reward: ~0.45–0.55.

3. Interactive Demo

python demo.py --episodes 3 --difficulty learning --delay 0.3

Watch the greedy agent handle a full shift in real time.

4. Diagnostics

python diagnose.py --episodes 10 --agent greedy

Per-episode breakdown: reward, fatalities, action histogram, events.

5. Visualize an Episode

# Export a greedy episode to JSON
python export_episode.py --difficulty learning --output episode.json

# Render as MP4 (pulsing calls, unit trails, UI overlay)
python visualize_env.py --input episode.json --output dispatchr_demo.mp4

Training

Prerequisites

  • Python 3.10+
  • CUDA-capable GPU (tested on L4 24GB and A100 80GB)
  • unsloth, transformers, torch, trl
pip install unsloth transformers torch trl numpy networkx matplotlib

Note: import unsloth must happen before import transformers to enable kernel optimizations.

Primary: TRL GRPO Training (Local)

python train_trl_grpo.py \
  --model Qwen/Qwen3-4B \
  --episodes 200 \
  --batch-size 4 \
  --num-generations 4 \
  --max-completion-length 256 \
  --curriculum \
  --save-every 25 \
  --output-dir ./outputs/trl_grpo

Why TRL-native (train_trl_grpo.py)?

  • Actively maintained training path in this repo
  • Random-step trajectory-cache dataset provides stronger learning signal
  • Integrated vLLM generation path for fast grouped rollouts on HF Jobs

Key flags:

  • --model Qwen/Qwen3-4B: Default stable model for TRL runs
  • --num-generations 4: Recommended minimum GRPO group size for non-zero within-group variance
  • --curriculum: Enable performance-gated difficulty escalation
  • --n-seeds: Number of cached trajectory seeds used for random-step sampling
  • --steps-per-seed: Number of random decision points sampled per seed
  • --trajectory-file: Saves every prompt/completion as JSONL for auditing
  • --push-to-hub --hub-model-id yourname/model: Push checkpoints to HF Hub automatically

Live Progress Tracking

The training script now includes built-in progress tracking:

  • Console reports after every batch: reward stats, moving averages, action rates, fatalities, ETA
  • JSONL log: outputs/trl_grpo/metrics.jsonl with per-batch data
  • Auto-generated plots at the end:
    • training_progress.png — 6-panel dashboard (reward, loss, validity, actions, fatalities, epsilon/curriculum)
    • training_curves.png — Minimal 2-panel (reward + loss)

Training on Hugging Face Hub Jobs (Recommended for Hackathon)

HF Hub Jobs provide serverless GPU compute — no Spaces needed, pay-as-you-go.

# 1. Install HF CLI
curl -LsSf https://hf.co/cli/install.sh | bash

# 2. Login
hf auth login

# 3. Launch training (TRL is default / main script)
python scripts/launch_hf_job.py --flavor l40sx1 --episodes 200

# Or A100 for faster throughput ($2.50/hr, 80GB)
python scripts/launch_hf_job.py \
  --flavor a100-large \
  --episodes 200

# 4. Watch logs stream in terminal
#    (Ctrl+C stops watching; job keeps running)

# 5. Check status anytime
hf jobs list
hf jobs logs <job-id>

Hardware Options:

Flavor GPU VRAM $/hr Best For
l4x1 L4 24 GB $0.80 4B model, budget runs, pilot testing
a100-large A100 80 GB $2.50 14B model, max performance, final runs

L4 vs A100:

  • L4 is better price/performance for 4B models (~$8-10 for 200 episodes)
  • A100 is required for 14B models (L4 OOMs at ~30GB loaded)
  • A100 is ~1.5–2× faster per batch but costs 3.1× more — only worth it for 14B

Budget math: $60 credit = 75 hrs on L4 or 24 hrs on A100.

Legacy Alternatives (Optional)

If you specifically want older training paths:

# Legacy manual GRPO
python train_grpo.py \
  --model Qwen/Qwen3.5-2B \
  --episodes 200 \
  --batch-size 8 \
  --use-4bit \
  --output-dir ./outputs/grpo

Same --curriculum flags supported.

# Legacy Unsloth path
python train_unsloth_grpo.py \
  --model unsloth/Qwen3-4B-Thinking-2507-bnb-4bit \
  --episodes 200 \
  --batch-size 8 \
  --output-dir ./outputs/unsloth_grpo

Monitoring Training Progress

During training (console):

tail -f outputs/unsloth_grpo/metrics.csv

After training (plots):

# Auto-generated at the end of train_unsloth_grpo.py
# outputs/unsloth_grpo/training_progress.png
# outputs/unsloth_grpo/training_curves.png

# Or generate manually from metrics.json
python scripts/plot_training.py --input outputs/unsloth_grpo/metrics.json

Track these metrics:

  • Mean reward — should increase from ~0.55 toward ~0.70+
  • MA5 / MA10 — moving averages smooth out episode variance
  • Loss — should be non-zero and gradually decrease
  • Valid action rate — should stay above 80%
  • Fatality count — should trend downward
  • Curriculum: phase=X — shows current difficulty phase
  • 🎓 CURRICULUM ESCALATED — confirms progression

Inspect trajectory:

python -c "
import json
with open('outputs/unsloth_grpo/trajectory.jsonl') as f:
    for line in f:
        d = json.loads(line)
        print(d['batch'], d['episode'], d['completion'][:100])
        break
"

Inference

Run a trained checkpoint:

# Standard loading
python inference.py --model-path ./outputs/unsloth_grpo/final --episodes 10

# Unsloth-optimized loading (faster)
python inference.py --model-path ./outputs/unsloth_grpo/final \
  --use-unsloth --load-in-4bit --episodes 10

# Save trajectory for visualization / before-after comparison
python inference.py --model-path ./outputs/unsloth_grpo/final \
  --episodes 3 --trajectory-file after.jsonl

Reward Design

The episode reward has 4 components:

  1. Response Score (50%): Ratio of oracle response time to actual response time. Capped at 1.0.
  2. Fatality Component (30%): 1.0 if zero deaths, penalized by 0.5 per fatality.
  3. Coverage Score (20%): Time-averaged fraction of steps where all 5 zones had a unit within 10 minutes.
  4. Action Validity Shaping: Small bonus for valid actions, penalty for invalid actions and excessive holding.
reward = 0.50 * response_score + 0.30 * fatality_component + 0.20 * coverage_score
         + validity_bonus - invalidity_penalty - idle_penalty

How the Environment Prevents Reward Hacking

Safeguard How It Works
Invalid action penalty Every invalid dispatch/stage/verify costs -0.02 in final reward
Idle penalty Every hold step costs -0.005; encourages active dispatching
Time-averaged coverage Coverage is measured per-step, not just at episode end. Can't cluster units at the end
Oracle uses start locations Oracle assignments are computed from unit start positions, not scattered end-state positions
Bridge collapse updates graph CityGraph recomputes all shortest paths when edges are destroyed. Units actually reroute
Radio delay is visible last_update_step shows when status was last confirmed, not current step
Ghost calls are detectable verify returns high confidence only 40% of the time on ghost calls (not 60% as design spec said)

OpenEnv Compliance

This environment implements the OpenEnv interface:

  • reset(difficulty) → initial observation
  • step(action_dict) → next observation, reward at episode end
  • stateState(episode_id, step_count)
  • get_ground_truth() → full episode metrics for analysis

Deploy to Hugging Face Spaces:

openenv push

Judging Metrics

After each episode, get_ground_truth() exposes:

{
  "fatality_count": 2,
  "valid_action_rate": 0.65,
  "invalid_action_rate": 0.1,
  "hold_rate": 0.25,
  "dispatch_rate": 0.4,
  "avg_response_time": 4.5,
  "calls_missed": 2,
  "response_score": 0.72,
  "coverage_score": 0.85
}

These metrics let judges verify that improvement is real, not just reward hacking.

Known Limitations

  • Training time: ~15–20 min per batch of 8 episodes on L4. 200 episodes ≈ 6–8 hours.
  • vLLM not used: Batched transformers.generate() is used instead because vLLM's KV cache conflicts with training memory on 24GB. Unsloth kernels mitigate the speed penalty.
  • Fire scene time longer: Fire calls tie up units for 5 steps (vs 3 for cardiac/trauma). This is correct but reduces effective fleet size during fire-heavy episodes.
  • Model occasionally echoes prompt: Rarely repeats system prompt text instead of generating an action when confused (mitigated by improved prompt with explicit examples).

Team

  • Team Name: TorchBearers
  • Team Members: Harsh Vardhan , Siddharth Kumar, Tejas Garg
  • Model: unsloth/Qwen3-4B-Thinking-2507-bnb-4bit (4B params, thinking-enabled, pre-quantized)
  • GPU: Lightning AI L4 (24GB VRAM) / HF Hub Jobs
  • Hackathon: Meta OpenEnv Hackathon, India, April 25-26 2026
  • Tests: 47/47 passing

About

No description, website, or topics provided.

Resources

Stars

2 stars

Watchers

0 watching

Forks

Releases

No releases published

Packages

 
 
 

Contributors