README.md

---
title: AI Crime Investigation World
emoji: 🔍
colorFrom: blue
colorTo: red
sdk: docker
app_port: 8000
---

# 🕵️ AI Crime Investigation World

> **Can an LLM learn to be a better detective?**
>
> This environment drops an AI agent into a crime scene with two suspects, a witness, and three pieces of physical evidence. The detective must interrogate, cross-reference alibis, catch contradictions, and make an accusation — all within 15 turns. Wrong guess? The criminal walks free.

🔗 **[Live Demo on Hugging Face Spaces](https://huggingface.co/spaces/printf-sourav/ai-crime-investigation-world)**

---

## The Problem

LLMs are great at answering questions, but terrible at **multi-step reasoning under uncertainty**. Real-world investigation requires:

- Asking the right questions to the right people
- Spotting when a story doesn't add up
- Weighing physical evidence against testimony
- Making a high-stakes decision with incomplete information

Current LLMs tend to accuse randomly, repeat the same questions, or ignore contradictions entirely. **Can RL fix that?**

## The Environment

Built on [OpenEnv](https://github.com/meta-pytorch/OpenEnv) for standardized RL training.

**What the agent sees:**
- A crime briefing (what happened, who's involved)
- Conversation history with suspects and witness
- Physical evidence log (keycard, CCTV, forensics)

**What the agent can do:**
- `ask_question` → Interrogate any suspect or the witness
- `request_evidence` → Examine one of three evidence types
- `accuse` → Name the criminal (episode ends)

**What ends an episode:**
- Accusation (correct = +10, wrong = −10)
- Timeout after 15 turns (−3 penalty)

**Why it's hard:**
- The guilty suspect has a convincing fake alibi
- The witness may be biased toward one suspect
- Evidence is ambiguous until cross-referenced with testimony
- Random early accusations are heavily penalized

## Hackathon Theme

**Theme #1: Multi-Agent Interactions** — Four AI agents with conflicting objectives (detective vs. guilty suspect vs. innocent suspect vs. biased witness) interact in a structured investigation. Each agent has private knowledge, separate reward signals, and behavioral incentives that create a rich strategic landscape.

## Reward Design (§7 — Multiple Independent Functions)

The detective's reward is shaped by **4 independent reward functions** during GRPO training and **11 event-based reward signals** during environment interaction:

| Training Reward Function | What It Measures |
|---|---|
| `format_reward_fn` | Is the output in valid `ACTION:` format? |
| `target_validity_reward_fn` | Does the action reference a valid agent/item? |
| `env_step_reward_fn` | Steps a fresh environment copy, returns actual delta |
| `strategic_reward_fn` | Rewards information gathering early, accusations late |

| Environment Event | Delta |
|---|---|
| Correct accusation | +10.0 |
| Wrong accusation | −10.0 (−12.0 with witness bias) |
| Timeout (no accusation) | −3.0 |
| Contradiction exposed | +2.0 |
| Prior-pattern exploited | +1.5 |
| Evidence confirms lead | +1.0 |
| Deflection resistance | +0.5 |
| Redundant question | −0.5 |
| Per-turn cost | −0.3 |

This combination encourages strategic questioning, contradiction detection, and evidence-backed accusations over random guessing.

## Anti-Reward-Hacking Safeguards (§8)

| Safeguard | How |
|---|---|
| Duplicate evidence blocked | `already_revealed` set tracks requested items |
| Redundant questions penalized | Per-topic tracking with `−0.5` penalty |
| Premature accusation blocked | `min_turns_before_accuse = 4` |
| Contradiction farming capped | Each topic rewarded at most once |
| Deflection bonus capped | Once per suspect |
| Timeout penalty | Forces decision within 15 turns |

## Curriculum Learning (§6)

Training starts easy and gets harder as the model improves:

| Difficulty | Turns | Evidence | Criminal | Witness Bias |
|---|---|---|---|---|
| Easy | 6 | 1 item | Always Suspect_A | None |
| Medium | 10 | 2 items | Random | None |
| Hard | 15 | 3 items | Random | Active |

Auto-advances when rolling accuracy ≥ 60% over last 10 evaluations.

## Training Stack (§10)

```
┌─────────────────────────────────────────────┐
│  TRL GRPOTrainer  ←  Policy optimization    │
│  Unsloth PatchFastRL  ←  2x speed           │
│  4-bit quantization (bitsandbytes)          │
│  LoRA (r=16, attention + MLP modules)       │
└─────────────────────────────────────────────┘
         ↕  reward_funcs
┌─────────────────────────────────────────────┐
│  CrimeInvestigationEnv (OpenEnv)            │
│  4 independent reward functions             │
│  Curriculum difficulty staging              │
└─────────────────────────────────────────────┘
```

### Quick Start — Training

```bash
# 1. Install PyTorch with CUDA (skip on Colab — pre-installed)
pip install torch --index-url https://download.pytorch.org/whl/cu126

# 2. Install Unsloth
pip install unsloth

# 3. Install remaining deps
pip install -r requirements-train.txt

# 4. Verify setup
python verify_setup.py

# 5. Train
python train_colab.py

# 6. Compare baseline vs trained
python eval_baseline.py
```

### Quick Start — Server

```bash
pip install -r requirements.txt
uvicorn server.app:app --host 0.0.0.0 --port 8000
# Open http://localhost:8000
```

## Results

<!-- TODO: Replace with actual training plots after running on Colab -->
_Training in progress — reward curves and before/after comparison will be added after the training run._

## Project Structure

```
crime_env/                      # Core environment package
  ├── case_generator.py             # Randomized crime scenario generation
  ├── environment.py                # step/reset/render RL interface
  ├── agent_prompts.py              # Role-specific system prompts
  ├── consistency_tracker.py        # Semantic contradiction detection
  ├── reward_calculator.py          # Multi-agent reward events
  └── constants.py                  # Shared constants
server/
  └── app.py                    # OpenEnv-compatible FastAPI server
dashboard.html                  # Interactive investigation dashboard
train_colab.py                  # TRL GRPO + Unsloth training script
eval_baseline.py                # Baseline vs trained model comparison
test_one_episode.py             # End-to-end test with scripted agents
verify_setup.py                 # Pre-training environment checker
openenv.yaml                    # OpenEnv manifest
Dockerfile                      # HuggingFace Spaces deployment
```

## API Endpoints (OpenEnv)

| Endpoint | Method | Description |
|----------|--------|-------------|
| `/` | GET | Interactive investigation dashboard |
| `/reset` | POST | Start a new episode (OpenEnv) |
| `/step` | POST | Execute a detective action (OpenEnv) |
| `/api/run_episode` | GET | Run full scripted episode, returns JSON trace |
| `/api/reward_curve` | GET | Reward history + smoothed metrics |
| `/api/health` | GET | Deployment health check |

## Links

- 🔗 **[Live Environment (HF Space)](https://huggingface.co/spaces/printf-sourav/ai-crime-investigation-world)**
- 📓 **Training Notebook**: _Coming soon (Colab link)_
- 📝 **Blog Post**: _Coming soon (HF blog link)_

## Known Limitations

- **Sparse terminal supervision**: The strongest reward signal is the terminal accusation outcome, so exploration quality matters early in GRPO training.
- **Rule-based NPC fallback**: The `_default_llm_call` uses string parsing to identify agent roles. Full LLM pipeline bypasses this.

## License

MIT