File size: 13,965 Bytes
fedfdcb
1c03487
 
 
 
 
 
 
 
 
3d5cf7f
1c03487
3d5cf7f
 
 
 
 
1c03487
3d5cf7f
1c03487
c2e4262
1c03487
c2e4262
 
3d5cf7f
c2e4262
3d5cf7f
 
 
 
c2e4262
3d5cf7f
c2e4262
 
 
3d5cf7f
c2e4262
3d5cf7f
fedfdcb
3d5cf7f
c2e4262
3d5cf7f
c2e4262
3d5cf7f
c2e4262
3d5cf7f
 
 
fedfdcb
3d5cf7f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
fedfdcb
3d5cf7f
 
 
 
 
 
 
 
bfbbcd1
3d5cf7f
6e28a1b
 
 
c2e4262
 
3d5cf7f
 
 
 
 
 
 
 
c2e4262
3d5cf7f
 
 
 
 
fedfdcb
6e28a1b
c2e4262
 
 
3d5cf7f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c2e4262
 
 
3d5cf7f
c2e4262
3d5cf7f
 
 
c2e4262
3d5cf7f
 
 
 
 
 
 
 
c2e4262
 
 
3d5cf7f
1c03487
3d5cf7f
 
6e28a1b
3d5cf7f
74f461a
c2e4262
6e28a1b
3d5cf7f
74f461a
3d5cf7f
6e28a1b
c2e4262
 
 
3d5cf7f
 
 
 
6e28a1b
 
3d5cf7f
 
 
c2e4262
6e28a1b
c2e4262
 
 
3d5cf7f
 
6e28a1b
c2e4262
3d5cf7f
 
dc89155
6e28a1b
74f461a
6e28a1b
3d5cf7f
6e28a1b
c2e4262
 
 
3d5cf7f
fedfdcb
6e28a1b
1c03487
3d5cf7f
fedfdcb
3d5cf7f
 
 
6e28a1b
3d5cf7f
 
6e28a1b
3d5cf7f
 
6e28a1b
3d5cf7f
219555f
3d5cf7f
c2e4262
1c03487
6e28a1b
fedfdcb
1e308ec
1c03487
6e28a1b
 
 
a251104
 
 
3d5cf7f
6e28a1b
1c03487
c2e4262
 
 
 
3d5cf7f
 
 
 
 
 
 
c2e4262
 
 
 
fedfdcb
 
3d5cf7f
 
 
 
 
219555f
c2e4262
6e28a1b
3d5cf7f
 
 
 
219555f
3d5cf7f
6e28a1b
3d5cf7f
 
 
c2e4262
3d5cf7f
219555f
3d5cf7f
 
6e28a1b
 
 
c2e4262
6e28a1b
3d5cf7f
6e28a1b
3d5cf7f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1e308ec
3d5cf7f
 
 
 
6e28a1b
3d5cf7f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
74f461a
3d5cf7f
c2e4262
 
 
1c03487
c2e4262
fedfdcb
3d5cf7f
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
---
title: Cascade Containment
emoji: 🦠
colorFrom: red
colorTo: blue
sdk: docker
app_port: 7860
pinned: false
---

## 🦠 An RL Benchmark for Sequential Resource Allocation Under Spreading Cascade Dynamics

[![OpenEnv](https://img.shields.io/badge/OpenEnv-Compliant-blue?style=flat-square)](https://github.com/meta-pytorch/OpenEnv)
[![Python](https://img.shields.io/badge/Python-3.10%2B-blue?style=flat-square)](https://python.org)
[![Docker](https://img.shields.io/badge/Docker-Ready-blue?style=flat-square)](https://hub.docker.com)
[![HF Space](https://img.shields.io/badge/HF%20Space-Live-green?style=flat-square)](https://huggingface.co/spaces/TheRubberDuckDebuggers/cascade-containment)
[![License](https://img.shields.io/badge/License-MIT-green?style=flat-square)](LICENSE)

Meta PyTorch OpenEnv Hackathon Γ— SST 2026 β€” [Live Demo](https://therubberduckdebuggers-cascade-containment.hf.space) Β· [GitHub](https://github.com/Rohitchandramouli/cascade-containment)

---

## The Problem

Sequential resource allocation under uncertainty is one of the most consequential decision problems in the real world. Whether containing an epidemic, deploying firefighting crews, isolating a cyberattack, or routing aid β€” the agent faces the same fundamental challenge:

- **Resources are scarce** β€” you cannot cover every district simultaneously
- **Data is delayed** β€” by the time a crisis is visible, it has already grown
- **Interventions cascade** β€” actions in one district affect adjacent ones
- **Acting too late is catastrophic** β€” hospital collapse ends the episode; proactive containment is rewarded exponentially more than reactive response

No existing OpenEnv benchmark formalises this problem class. Cascade Containment does.

---

## Environment Overview

A city health authority must allocate limited medical resources across districts to contain a spreading outbreak. Each step, the agent observes district infection rates (possibly lagged), hospital capacity levels, and growth signals β€” then decides where to deploy resources, impose restrictions, or gather data.

The environment is **not epidemic-specific**. The underlying mechanics β€” spreading cascade, delayed observation, resource scarcity, geographic spillover β€” are structurally identical across multiple real-world domains.

---

## Quick Start

```python
from client import CascadeContainmentEnv
from models import ContainmentAction

with CascadeContainmentEnv(
    base_url="https://therubberduckdebuggers-cascade-containment.hf.space"
).sync() as env:
    result = env.reset(task_name="medium")
    obs    = result.observation

    while not result.done:
        most_infected = max(obs.districts, key=lambda d: d.reported_infection_rate)
        action = ContainmentAction(
            action_type="allocate",
            district_id=most_infected.district_id
        )
        result = env.step(action)
        obs    = result.observation
        print(f"Step {obs.current_step}: reward={result.reward:.3f}")
```

## Running the Full Baseline Evaluation

```bash
export API_BASE_URL="https://router.huggingface.co/v1"
export MODEL_NAME="meta-llama/Llama-3.1-8B-Instruct"
export HF_TOKEN="hf_your_token_here"
export ENV_BASE_URL="https://therubberduckdebuggers-cascade-containment.hf.space"

# Full LLM+GRPO evaluation (~10 minutes, emits [START][STEP][END] logs)
python inference.py

# Local validation + greedy benchmark
python scripts/test_local.py
```

---

## Action Space

| Field | Type | Values |
| --- | --- | --- |
| `action_type` | `string` | `"test"` Β· `"restrict"` Β· `"allocate"` |
| `district_id` | `int` | 0-indexed district target |

| Action | Cost | Effect |
| --- | --- | --- |
| **test** | 1 resource | Reveals accurate current infection data for district |
| **restrict** | Free | Imposes movement restrictions; reduces spread rate; penalised if infection < 0.20 |
| **allocate** | 1 resource | Deploys medical resources; reduces existing infection by 5% and slows future spread |

Movement restrictions lift automatically once a district's infection drops below the safe threshold.

---

## Observation Space

The agent receives a filtered, potentially lagged view of the world β€” **never the full ground truth**:

```python
CityObservation:
  districts:            List[DistrictObservation]  # per-district visible state
  available_resources:  int                         # budget remaining this step
  current_step:         int
  max_steps:            int
  done:                 bool
  reward:               float | None
  message:              str | None
```

| Field | Description | Observability |
| --- | --- | --- |
| `reported_infection_rate` | Active infection fraction | Real-time (easy/medium); **3-day lagged** (hard) |
| `growth_rate_hint` | Noisy signal of true spread rate | Always real-time Β± noise |
| `hospital_capacity_remaining` | ICU/ward capacity fraction | Always real-time |
| `population_density` | District's share of city population | Always real-time |
| `restriction_active` | Whether movement restrictions are in place | Always real-time |
| `tested_recently` | Tested within last 2 days | Always real-time |

---

## Epidemiological Model

```text
new_infection = current + (spread_rate βˆ’ natural_recovery βˆ’ intervention) + geographic_spillover
```

| Parameter | Value | Rationale |
| --- | --- | --- |
| Spread rate | 3–8% per day | Realistic for respiratory outbreaks (seasonal flu: 5–10%) |
| Natural recovery | 1% per day | Background case resolution without medical intervention |
| Treatment effect | βˆ’5% existing infection | Medical deployment (antivirals, PPE, rapid response) |
| Spread reduction | βˆ’10% per allocation | Resource-driven suppression of transmission |
| Geographic spillover | 1% to adjacent districts | Linear topology β€” no wrap-around (geographically realistic) |
| Hospital breach threshold | ≀10% capacity | Real ICU overflow and triage failure threshold |

---

## Three Tasks

| Task | Districts | Steps | Resources | Data Lag | Challenge |
| --- | --- | --- | --- | --- | --- |
| **easy** | 2 | 10 | 10 | None | Single outbreak; D1 starts infected, D0 is clean |
| **medium** | 4 | 15 | 8 | None | Two simultaneous outbreaks; forced triage between competing threats |
| **hard** | 6 | 15 | 7 | **3 days** | Six seeded infections (only D2 and D4 above safe threshold); 3-day data lag; scarce resources |

**Easy** β€” D1 starts at 0.50 infection, D0 is clean. The agent must observe and target the correct district. A fixed-target agent ignoring observations scores ~43% and breaches hospitals 60% of the time.

**Medium** β€” D0 starts above the infection threshold (0.42); D2 is in the warning zone (0.38, below the 0.40 threshold). D1 and D3 start low but grow into crisis within 4–6 steps via spillover. With 8 resources across 4 districts over 15 steps, genuine triage is required.

**Hard** β€” 3-day information lag means the agent sees infection rates from 3 days ago. The `growth_rate_hint` provides a noisy signal to estimate current state. Structural uncertainty β€” not testable around.

---

## Reward Function

| Term | Value | Fires When |
| --- | --- | --- |
| Infection penalty | `βˆ’0.50 Γ— density` | District infection > 0.40 |
| Hospital breach | `βˆ’1.00` | Hospital capacity ≀ 10% |
| Early containment | `+0.50 Γ— (1 βˆ’ step/max_steps)` | District infection < 0.20 |
| Correct prioritisation | `+0.30` | Allocate to highest-infected district |
| Unnecessary restriction | `βˆ’0.20` | Restrict district below 0.20 |

Key design choices: early containment bonus decays over time (proactive action worth more); infection penalty scales with population density (realistic triage); hospital breach fires at 10% not 0% (real ICU thresholds); restrictions auto-lift when infection drops below safe threshold.

---

## Grader

Fully deterministic β€” no randomness, no LLM calls. Identical trajectories always produce identical scores in `[0.0, 1.0]`.

| Component | Weight | Measures |
| --- | --- | --- |
| **Hospital score** | 45% | Average capacity preserved; Γ—0.6 multiplier if any district collapsed |
| **Containment score** | 30% | Fraction of district-days below infection threshold (first 2 steps excluded) |
| **Efficiency score** | 15% | Fraction of resource actions targeting highest-infected district (grader uses pre-action state to avoid penalising successful treatments) |
| **Speed score** | 10% | `1 βˆ’ (steps / max_steps)` if episode ends early; else 0 |

Hospital is weighted highest because system capacity preservation is the primary operational constraint in real outbreak response β€” a functioning healthcare system is the prerequisite for everything else.

---

## Baseline Agent β€” GRPO-Style Episodic Memory

No weight updates, no gradient computation. The prompt is the policy; memory updates are the policy improvement.

### Learning Loop

```text
Rollout 1: Base prompt, no prior knowledge
  compute advantage = R1 - mean([])
  store steps with reward > -0.3 into EpisodicMemory

Rollout 2: Memory-augmented prompt
  retrieve top-5 similar past decisions by L1 distance on infection profiles
  inject as concrete examples into prompt
  compute advantage = R2 - mean([R1])
  reinforce if advantage > -0.5

... repeat for N rollouts (easy=2, medium=3, hard=3)
Report best grader score across all rollouts
```

### Benchmark Scores

Results from `baseline/run.py` (Llama 3.1 8B Instant via Groq, runtime ~10 minutes):

| Task | Greedy (D0) | LLM+GRPO | Lift |
| --- | --- | --- | --- |
| Easy | 42.8% (breach 60%) | **90.8%** | +48pp |
| Medium | 39.6% (breach 100%) | **78.0%** | +38pp |
| Hard | 35.3% (breach 100%) | **61.1%** | +26pp |
| **Average** | **39.2%** | **75.9%** | **+37pp** |

The greedy baseline (always allocates to D0) scores 33–43% with 60–100% hospital breach rates β€” no trivial exploit path. LLM+GRPO reaches 66–91% with zero breaches, demonstrating that genuine triage reasoning is required and rewarded.

---

## Generalisation

| Domain | Spreading cascade | Delayed data | Resource scarcity |
| --- | --- | --- | --- |
| 🦠 **Epidemic containment** | Infection spreads between districts | Lagged case counts | Medical resources |
| πŸ”₯ **Wildfire deployment** | Fire spreads across terrain | Satellite update delay | Firefighting crews |
| πŸ›‘οΈ **Cyberattack isolation** | Lateral movement between systems | Detection lag | Security team hours |
| πŸ“’ **Misinformation containment** | Narrative spread through networks | Viral detection lag | Correction budget |
| 🀝 **Poverty intervention** | Deprivation cascades through communities | Census data lag | Aid allocation |

---

## Project Structure

```text
cascade-containment/
β”œβ”€β”€ inference.py              # Evaluation entry point (mandatory [START][STEP][END] logs)
β”œβ”€β”€ models.py                 # Typed data contracts: Action, Observation, State
β”œβ”€β”€ client.py                 # OpenEnv client interface
β”œβ”€β”€ openenv.yaml              # Environment manifest for OpenEnv registry
β”‚
β”œβ”€β”€ server/
β”‚   β”œβ”€β”€ app.py                # FastAPI server + judge dashboard + /grade /info /demo /validate endpoints
β”‚   β”œβ”€β”€ environment.py        # Core RL loop (reset/step/state OpenEnv interface)
β”‚   β”œβ”€β”€ grader.py             # Deterministic trajectory scorer β€” no LLM calls
β”‚   β”œβ”€β”€ constants.py          # Single source of truth for all numeric configuration
β”‚   β”œβ”€β”€ utils.py              # Spread computation, observation builder, helper functions
β”‚   β”œβ”€β”€ Dockerfile            # Container definition
β”‚   └── tasks/
β”‚       β”œβ”€β”€ task_easy.py      # 2 districts, 10 steps, real-time data (D1 seeded)
β”‚       β”œβ”€β”€ task_medium.py    # 4 districts, 15 steps, forced triage
β”‚       └── task_hard.py      # 6 districts, 15 steps, 3-day data lag
β”‚
β”œβ”€β”€ baseline/
β”‚   β”œβ”€β”€ policy.py             # LLM policy with chain-of-thought prompting
β”‚   β”œβ”€β”€ evaluator.py          # GRPO episodic memory loop (easy=2, medium=3, hard=3 rollouts)
β”‚   └── run.py                # CLI entry point
β”‚
β”œβ”€β”€ scripts/
β”‚   └── test_local.py         # Phase 1 spec checks + greedy benchmark + variance analysis
β”‚
└── core/
    β”œβ”€β”€ trajectory.py         # EpisodicMemory β€” L1 similarity retrieval, phase-weighted
    β”œβ”€β”€ reward.py             # Score normalisation utilities
    └── policy_update.py      # Advantage computation, memory gating (threshold -0.3)
```

---

## OpenEnv Compliance

| Requirement | Status |
| --- | --- |
| `reset()` returns `CityObservation` | βœ… |
| `step(action)` returns `CityObservation` | βœ… |
| `state` property returns `State` | βœ… |
| Typed `Action` subclass | βœ… `ContainmentAction(Action)` |
| Typed `Observation` subclass | βœ… `CityObservation(Observation)` |
| `openenv.yaml` manifest | βœ… |
| Dockerfile builds | βœ… |
| HF Space deploys | βœ… |
| `inference.py` at root | βœ… |
| `[START][STEP][END]` structured logs | βœ… |
| Runtime < 20 minutes | βœ… ~~10 minutes |
| `API_BASE_URL`, `MODEL_NAME`, `HF_TOKEN` env vars | βœ… |
| OpenAI client for all LLM calls | βœ… |
| Grader scores in `[0.0, 1.0]` | βœ… |
| 3+ tasks with difficulty progression | βœ… |
| Phase 1 automated validation | βœ… 10/10 checks pass |

---

## Setup and Local Development

### Local Server

```bash
pip install -r requirements.txt
uvicorn server.app:app --host 0.0.0.0 --port 7860

export ENV_BASE_URL=http://localhost:7860
python baseline/run.py
```

### Docker

```bash
docker build -t cascade-containment .
docker run -p 7860:7860 cascade-containment
```

---

## Tags

`reinforcement-learning` Β· `resource-allocation` Β· `sequential-decision-making` Β· `partial-observability` Β· `cascade-dynamics` Β· `epidemic-response` Β· `openenv` Β· `llm-agent` Β· `grpo` Β· `episodic-memory` Β· `triage` Β· `multi-district` Β· `docker` Β· `fastapi`