Spaces:
Sleeping
Sleeping
Upload folder using huggingface_hub
Browse filesThis view is limited to 50 files because it contains too many changes. See raw diff
- CLAUDE.md +4 -4
- README.md +255 -255
- __init__.py +16 -16
- cortex/brains/__init__.py +26 -0
- cortex/brains/_base.py +201 -0
- cortex/brains/_executive.py +144 -0
- cortex/brains/epidemiology.py +25 -0
- cortex/brains/governance.py +22 -0
- cortex/brains/logistics.py +23 -0
- cortex/lenses.py +186 -0
- cortex/schemas.py +46 -1
- cortex/subagents/__init__.py +27 -0
- cortex/subagents/_base.py +240 -0
- cortex/subagents/critic.py +68 -0
- cortex/subagents/perception.py +192 -0
- cortex/subagents/planner.py +61 -0
- cortex/subagents/prompts/critic.txt +19 -0
- cortex/subagents/prompts/planner.txt +26 -0
- cortex/subagents/prompts/world_modeler.txt +32 -0
- cortex/subagents/world_modeler.py +61 -0
- demo/CLAUDE.md +45 -45
- docs/CORTEX_ARCHITECTURE.md +1 -1
- inference.py +9 -4
- notebooks/train_b1_grpo.ipynb +496 -0
- notebooks/train_cortex_router_grpo.ipynb +567 -0
- openenv.yaml +7 -7
- openenv_CrisisWorldCortex.egg-info/PKG-INFO +13 -0
- openenv_CrisisWorldCortex.egg-info/SOURCES.txt +48 -0
- openenv_CrisisWorldCortex.egg-info/dependency_links.txt +1 -0
- openenv_CrisisWorldCortex.egg-info/entry_points.txt +2 -0
- openenv_CrisisWorldCortex.egg-info/requires.txt +9 -0
- openenv_CrisisWorldCortex.egg-info/top_level.txt +1 -0
- pyproject.toml +1 -0
- server/CLAUDE.md +9 -18
- server/CrisisWorldCortex_environment.py +19 -3
- server/__init__.py +11 -11
- server/app.py +2 -2
- server/graders/outer_reward.py +96 -34
- server/requirements.txt +6 -6
- tests/CLAUDE.md +48 -48
- tests/_helpers/__init__.py +6 -0
- tests/_helpers/llm_stub.py +82 -0
- tests/test_baseline_b1.py +28 -30
- tests/test_cortex_brain_executive.py +206 -0
- tests/test_cortex_brain_smoke.py +142 -0
- tests/test_cortex_lenses.py +238 -0
- tests/test_cortex_perception.py +116 -0
- tests/test_cortex_subagents.py +387 -0
- tests/test_outer_reward_in_range.py +12 -10
- tests/test_outer_reward_non_constancy.py +14 -5
CLAUDE.md
CHANGED
|
@@ -28,7 +28,7 @@ Wire-protocol classes are template-locked: `CrisisworldcortexAction`, `Crisiswor
|
|
| 28 |
|
| 29 |
- The four wire-protocol class names above. Do not rename.
|
| 30 |
- Flat package layout: `pyproject.toml:package-dir` maps `CrisisWorldCortex` to repo root. Do not move `models.py`, `client.py`, or root `__init__.py` off root.
|
| 31 |
-
-
|
| 32 |
- `pyproject.toml`: deps, entry points, `packages`, `package-dir`.
|
| 33 |
- `openenv.yaml` keys: `spec_version`, `runtime`, `app`, `port`.
|
| 34 |
- `server/app.py:create_app(...)` call signature.
|
|
@@ -47,7 +47,7 @@ Wire-protocol classes are template-locked: `CrisisworldcortexAction`, `Crisiswor
|
|
| 47 |
## Import-graph rule (enforced)
|
| 48 |
|
| 49 |
- `cortex/` imports `models` and `cortex/*` only.
|
| 50 |
-
- `server/` imports `models`, `openenv.core.*`,
|
| 51 |
- `baselines/` imports `models`, `client`, `cortex/*`. Must not import `server/*` — baselines hit the env over HTTP like production.
|
| 52 |
- `training/` imports `models`, `client`, `cortex/*`, `server.graders` (reward-name constants only). Must not import `server.simulator/*`.
|
| 53 |
- `demo/` imports `cortex.schemas` (types only) and stdlib. Must not import `server/*`, `training/*`, `baselines/*`, `cortex.council`, `cortex.routing_policy`.
|
|
@@ -70,9 +70,9 @@ Do not restate subsystem APIs in this file or in other subsystem files.
|
|
| 70 |
|
| 71 |
- Wire package: `from CrisisWorldCortex import CrisisworldcortexAction, ...`.
|
| 72 |
- Dev / research directories (sibling-of-repo-root): bare-name — `import cortex`, `import baselines`, `import training`, `import demo`, `import scripts`.
|
| 73 |
-
- Server-internal:
|
| 74 |
- **Cross-boundary into wire types**: when a dev / research directory (`cortex/`, `baselines/`, `training/`, `demo/`) crosses into the wire-protocol package, use `from CrisisWorldCortex.models import …` — **never** bare `from models import …`. Bare-name only applies to dev-directory siblings and to same-subpackage imports inside `server/`. The dual-path import creates distinct `sys.modules` entries and breaks Pydantic discriminator checks across import boundaries (verified by session 4's class-identity bug; `cortex/schemas.py:21` documents the canonicalised import).
|
| 75 |
-
- **Cross-boundary into wire types from server modules**:
|
| 76 |
|
| 77 |
## Commands (Git Bash; quote Windows paths)
|
| 78 |
|
|
|
|
| 28 |
|
| 29 |
- The four wire-protocol class names above. Do not rename.
|
| 30 |
- Flat package layout: `pyproject.toml:package-dir` maps `CrisisWorldCortex` to repo root. Do not move `models.py`, `client.py`, or root `__init__.py` off root.
|
| 31 |
+
- Dual-import fallback pattern (`try: from ..models / except: from models`) in every `server/` module that imports `models`.
|
| 32 |
- `pyproject.toml`: deps, entry points, `packages`, `package-dir`.
|
| 33 |
- `openenv.yaml` keys: `spec_version`, `runtime`, `app`, `port`.
|
| 34 |
- `server/app.py:create_app(...)` call signature.
|
|
|
|
| 47 |
## Import-graph rule (enforced)
|
| 48 |
|
| 49 |
- `cortex/` imports `models` and `cortex/*` only.
|
| 50 |
+
- `server/` imports `models`, `openenv.core.*`, `server/*` only. No `cortex/*`, no `training/*`, no `baselines/*`, no `demo/*`.
|
| 51 |
- `baselines/` imports `models`, `client`, `cortex/*`. Must not import `server/*` — baselines hit the env over HTTP like production.
|
| 52 |
- `training/` imports `models`, `client`, `cortex/*`, `server.graders` (reward-name constants only). Must not import `server.simulator/*`.
|
| 53 |
- `demo/` imports `cortex.schemas` (types only) and stdlib. Must not import `server/*`, `training/*`, `baselines/*`, `cortex.council`, `cortex.routing_policy`.
|
|
|
|
| 70 |
|
| 71 |
- Wire package: `from CrisisWorldCortex import CrisisworldcortexAction, ...`.
|
| 72 |
- Dev / research directories (sibling-of-repo-root): bare-name — `import cortex`, `import baselines`, `import training`, `import demo`, `import scripts`.
|
| 73 |
+
- Server-internal: `from server.simulator import ...`, `from server.graders import ...`.
|
| 74 |
- **Cross-boundary into wire types**: when a dev / research directory (`cortex/`, `baselines/`, `training/`, `demo/`) crosses into the wire-protocol package, use `from CrisisWorldCortex.models import …` — **never** bare `from models import …`. Bare-name only applies to dev-directory siblings and to same-subpackage imports inside `server/`. The dual-path import creates distinct `sys.modules` entries and breaks Pydantic discriminator checks across import boundaries (verified by session 4's class-identity bug; `cortex/schemas.py:21` documents the canonicalised import).
|
| 75 |
+
- **Cross-boundary into wire types from deep server modules**: files inside `server/` more than one level deep (e.g., `server/simulator/seir_model.py`, `server/graders/outer_reward.py`) cannot use the `try: from ..models / except: from models` fallback — `..models` from a two-level-deep module resolves to a non-existent `CrisisWorldCortex.server.models`, the fallback fires, and bare `models` loads as a separate `sys.modules` entry. Use the absolute path: `from CrisisWorldCortex.models import …`. Session 4 (`cortex/schemas.py:21`) and Session 5a (`server/simulator/seir_model.py`, `server/simulator/tasks.py`) document this with inline comments. The dual-import fallback in `server/CrisisWorldCortex_environment.py` and `server/app.py` works only because they are one level deep (`..models` → `CrisisWorldCortex.models` directly).
|
| 76 |
|
| 77 |
## Commands (Git Bash; quote Windows paths)
|
| 78 |
|
README.md
CHANGED
|
@@ -1,255 +1,255 @@
|
|
| 1 |
-
---
|
| 2 |
-
title: Crisisworldcortex Environment Server
|
| 3 |
-
emoji: 🌟
|
| 4 |
-
colorFrom: yellow
|
| 5 |
-
colorTo: pink
|
| 6 |
-
sdk: docker
|
| 7 |
-
pinned: false
|
| 8 |
-
app_port: 8000
|
| 9 |
-
base_path: /web
|
| 10 |
-
tags:
|
| 11 |
-
- openenv
|
| 12 |
-
---
|
| 13 |
-
|
| 14 |
-
# Crisisworldcortex Environment
|
| 15 |
-
|
| 16 |
-
A simple test environment that echoes back messages. Perfect for testing the env APIs as well as demonstrating environment usage patterns.
|
| 17 |
-
|
| 18 |
-
## Quick Start
|
| 19 |
-
|
| 20 |
-
The simplest way to use the Crisisworldcortex environment is through the `CrisisworldcortexEnv` class:
|
| 21 |
-
|
| 22 |
-
```python
|
| 23 |
-
from CrisisWorldCortex import CrisisworldcortexAction, CrisisworldcortexEnv
|
| 24 |
-
|
| 25 |
-
try:
|
| 26 |
-
# Create environment from Docker image
|
| 27 |
-
CrisisWorldCortexenv = CrisisworldcortexEnv.from_docker_image("CrisisWorldCortex-env:latest")
|
| 28 |
-
|
| 29 |
-
# Reset
|
| 30 |
-
result = CrisisWorldCortexenv.reset()
|
| 31 |
-
print(f"Reset: {result.observation.echoed_message}")
|
| 32 |
-
|
| 33 |
-
# Send multiple messages
|
| 34 |
-
messages = ["Hello, World!", "Testing echo", "Final message"]
|
| 35 |
-
|
| 36 |
-
for msg in messages:
|
| 37 |
-
result = CrisisWorldCortexenv.step(CrisisworldcortexAction(message=msg))
|
| 38 |
-
print(f"Sent: '{msg}'")
|
| 39 |
-
print(f" → Echoed: '{result.observation.echoed_message}'")
|
| 40 |
-
print(f" → Length: {result.observation.message_length}")
|
| 41 |
-
print(f" → Reward: {result.reward}")
|
| 42 |
-
|
| 43 |
-
finally:
|
| 44 |
-
# Always clean up
|
| 45 |
-
CrisisWorldCortexenv.close()
|
| 46 |
-
```
|
| 47 |
-
|
| 48 |
-
That's it! The `CrisisworldcortexEnv.from_docker_image()` method handles:
|
| 49 |
-
- Starting the Docker container
|
| 50 |
-
- Waiting for the server to be ready
|
| 51 |
-
- Connecting to the environment
|
| 52 |
-
- Container cleanup when you call `close()`
|
| 53 |
-
|
| 54 |
-
## Building the Docker Image
|
| 55 |
-
|
| 56 |
-
Before using the environment, you need to build the Docker image:
|
| 57 |
-
|
| 58 |
-
```bash
|
| 59 |
-
# From project root
|
| 60 |
-
docker build -t CrisisWorldCortex-env:latest -f server/Dockerfile .
|
| 61 |
-
```
|
| 62 |
-
|
| 63 |
-
## Deploying to Hugging Face Spaces
|
| 64 |
-
|
| 65 |
-
You can easily deploy your OpenEnv environment to Hugging Face Spaces using the `openenv push` command:
|
| 66 |
-
|
| 67 |
-
```bash
|
| 68 |
-
# From the environment directory (where openenv.yaml is located)
|
| 69 |
-
openenv push
|
| 70 |
-
|
| 71 |
-
# Or specify options
|
| 72 |
-
openenv push --namespace my-org --private
|
| 73 |
-
```
|
| 74 |
-
|
| 75 |
-
The `openenv push` command will:
|
| 76 |
-
1. Validate that the directory is an OpenEnv environment (checks for `openenv.yaml`)
|
| 77 |
-
2. Prepare a custom build for Hugging Face Docker space (enables web interface)
|
| 78 |
-
3. Upload to Hugging Face (ensuring you're logged in)
|
| 79 |
-
|
| 80 |
-
### Prerequisites
|
| 81 |
-
|
| 82 |
-
- Authenticate with Hugging Face: The command will prompt for login if not already authenticated
|
| 83 |
-
|
| 84 |
-
### Options
|
| 85 |
-
|
| 86 |
-
- `--directory`, `-d`: Directory containing the OpenEnv environment (defaults to current directory)
|
| 87 |
-
- `--repo-id`, `-r`: Repository ID in format 'username/repo-name' (defaults to 'username/env-name' from openenv.yaml)
|
| 88 |
-
- `--base-image`, `-b`: Base Docker image to use (overrides Dockerfile FROM)
|
| 89 |
-
- `--private`: Deploy the space as private (default: public)
|
| 90 |
-
|
| 91 |
-
### Examples
|
| 92 |
-
|
| 93 |
-
```bash
|
| 94 |
-
# Push to your personal namespace (defaults to username/env-name from openenv.yaml)
|
| 95 |
-
openenv push
|
| 96 |
-
|
| 97 |
-
# Push to a specific repository
|
| 98 |
-
openenv push --repo-id my-org/my-env
|
| 99 |
-
|
| 100 |
-
# Push with a custom base image
|
| 101 |
-
openenv push --base-image ghcr.io/meta-pytorch/openenv-base:latest
|
| 102 |
-
|
| 103 |
-
# Push as a private space
|
| 104 |
-
openenv push --private
|
| 105 |
-
|
| 106 |
-
# Combine options
|
| 107 |
-
openenv push --repo-id my-org/my-env --base-image custom-base:latest --private
|
| 108 |
-
```
|
| 109 |
-
|
| 110 |
-
After deployment, your space will be available at:
|
| 111 |
-
`https://huggingface.co/spaces/<repo-id>`
|
| 112 |
-
|
| 113 |
-
The deployed space includes:
|
| 114 |
-
- **Web Interface** at `/web` - Interactive UI for exploring the environment
|
| 115 |
-
- **API Documentation** at `/docs` - Full OpenAPI/Swagger interface
|
| 116 |
-
- **Health Check** at `/health` - Container health monitoring
|
| 117 |
-
- **WebSocket** at `/ws` - Persistent session endpoint for low-latency interactions
|
| 118 |
-
|
| 119 |
-
## Environment Details
|
| 120 |
-
|
| 121 |
-
### Action
|
| 122 |
-
**CrisisworldcortexAction**: Contains a single field
|
| 123 |
-
- `message` (str) - The message to echo back
|
| 124 |
-
|
| 125 |
-
### Observation
|
| 126 |
-
**CrisisworldcortexObservation**: Contains the echo response and metadata
|
| 127 |
-
- `echoed_message` (str) - The message echoed back
|
| 128 |
-
- `message_length` (int) - Length of the message
|
| 129 |
-
- `reward` (float) - Reward based on message length (length × 0.1)
|
| 130 |
-
- `done` (bool) - Always False for echo environment
|
| 131 |
-
- `metadata` (dict) - Additional info like step count
|
| 132 |
-
|
| 133 |
-
### Reward
|
| 134 |
-
The reward is calculated as: `message_length × 0.1`
|
| 135 |
-
- "Hi" → reward: 0.2
|
| 136 |
-
- "Hello, World!" → reward: 1.3
|
| 137 |
-
- Empty message → reward: 0.0
|
| 138 |
-
|
| 139 |
-
## Advanced Usage
|
| 140 |
-
|
| 141 |
-
### Connecting to an Existing Server
|
| 142 |
-
|
| 143 |
-
If you already have a Crisisworldcortex environment server running, you can connect directly:
|
| 144 |
-
|
| 145 |
-
```python
|
| 146 |
-
from CrisisWorldCortex import CrisisworldcortexEnv
|
| 147 |
-
|
| 148 |
-
# Connect to existing server
|
| 149 |
-
CrisisWorldCortexenv = CrisisworldcortexEnv(base_url="<ENV_HTTP_URL_HERE>")
|
| 150 |
-
|
| 151 |
-
# Use as normal
|
| 152 |
-
result = CrisisWorldCortexenv.reset()
|
| 153 |
-
result = CrisisWorldCortexenv.step(CrisisworldcortexAction(message="Hello!"))
|
| 154 |
-
```
|
| 155 |
-
|
| 156 |
-
Note: When connecting to an existing server, `CrisisWorldCortexenv.close()` will NOT stop the server.
|
| 157 |
-
|
| 158 |
-
### Using the Context Manager
|
| 159 |
-
|
| 160 |
-
The client supports context manager usage for automatic connection management:
|
| 161 |
-
|
| 162 |
-
```python
|
| 163 |
-
from CrisisWorldCortex import CrisisworldcortexAction, CrisisworldcortexEnv
|
| 164 |
-
|
| 165 |
-
# Connect with context manager (auto-connects and closes)
|
| 166 |
-
with CrisisworldcortexEnv(base_url="http://localhost:8000") as env:
|
| 167 |
-
result = env.reset()
|
| 168 |
-
print(f"Reset: {result.observation.echoed_message}")
|
| 169 |
-
# Multiple steps with low latency
|
| 170 |
-
for msg in ["Hello", "World", "!"]:
|
| 171 |
-
result = env.step(CrisisworldcortexAction(message=msg))
|
| 172 |
-
print(f"Echoed: {result.observation.echoed_message}")
|
| 173 |
-
```
|
| 174 |
-
|
| 175 |
-
The client uses WebSocket connections for:
|
| 176 |
-
- **Lower latency**: No HTTP connection overhead per request
|
| 177 |
-
- **Persistent session**: Server maintains your environment state
|
| 178 |
-
- **Efficient for episodes**: Better for many sequential steps
|
| 179 |
-
|
| 180 |
-
### Concurrent WebSocket Sessions
|
| 181 |
-
|
| 182 |
-
The server supports multiple concurrent WebSocket connections. To enable this,
|
| 183 |
-
modify `server/app.py` to use factory mode:
|
| 184 |
-
|
| 185 |
-
```python
|
| 186 |
-
# In server/app.py - use factory mode for concurrent sessions
|
| 187 |
-
app = create_app(
|
| 188 |
-
CrisisworldcortexEnvironment, # Pass class, not instance
|
| 189 |
-
CrisisworldcortexAction,
|
| 190 |
-
CrisisworldcortexObservation,
|
| 191 |
-
max_concurrent_envs=4, # Allow 4 concurrent sessions
|
| 192 |
-
)
|
| 193 |
-
```
|
| 194 |
-
|
| 195 |
-
Then multiple clients can connect simultaneously:
|
| 196 |
-
|
| 197 |
-
```python
|
| 198 |
-
from CrisisWorldCortex import CrisisworldcortexAction, CrisisworldcortexEnv
|
| 199 |
-
from concurrent.futures import ThreadPoolExecutor
|
| 200 |
-
|
| 201 |
-
def run_episode(client_id: int):
|
| 202 |
-
with CrisisworldcortexEnv(base_url="http://localhost:8000") as env:
|
| 203 |
-
result = env.reset()
|
| 204 |
-
for i in range(10):
|
| 205 |
-
result = env.step(CrisisworldcortexAction(message=f"Client {client_id}, step {i}"))
|
| 206 |
-
return client_id, result.observation.message_length
|
| 207 |
-
|
| 208 |
-
# Run 4 episodes concurrently
|
| 209 |
-
with ThreadPoolExecutor(max_workers=4) as executor:
|
| 210 |
-
results = list(executor.map(run_episode, range(4)))
|
| 211 |
-
```
|
| 212 |
-
|
| 213 |
-
## Development & Testing
|
| 214 |
-
|
| 215 |
-
### Direct Environment Testing
|
| 216 |
-
|
| 217 |
-
Test the environment logic directly without starting the HTTP server:
|
| 218 |
-
|
| 219 |
-
```bash
|
| 220 |
-
# From the server directory
|
| 221 |
-
python3 server/CrisisWorldCortex_environment.py
|
| 222 |
-
```
|
| 223 |
-
|
| 224 |
-
This verifies that:
|
| 225 |
-
- Environment resets correctly
|
| 226 |
-
- Step executes actions properly
|
| 227 |
-
- State tracking works
|
| 228 |
-
- Rewards are calculated correctly
|
| 229 |
-
|
| 230 |
-
### Running Locally
|
| 231 |
-
|
| 232 |
-
Run the server locally for development:
|
| 233 |
-
|
| 234 |
-
```bash
|
| 235 |
-
uvicorn server.app:app --reload
|
| 236 |
-
```
|
| 237 |
-
|
| 238 |
-
## Project Structure
|
| 239 |
-
|
| 240 |
-
```
|
| 241 |
-
CrisisWorldCortex/
|
| 242 |
-
├── .dockerignore # Docker build exclusions
|
| 243 |
-
├── __init__.py # Module exports
|
| 244 |
-
├── README.md # This file
|
| 245 |
-
├── openenv.yaml # OpenEnv manifest
|
| 246 |
-
├── pyproject.toml # Project metadata and dependencies
|
| 247 |
-
├── uv.lock # Locked dependencies (generated)
|
| 248 |
-
├── client.py # CrisisworldcortexEnv client
|
| 249 |
-
├── models.py # Action and Observation models
|
| 250 |
-
└── server/
|
| 251 |
-
├── __init__.py # Server module exports
|
| 252 |
-
├── CrisisWorldCortex_environment.py # Core environment logic
|
| 253 |
-
├── app.py # FastAPI application (HTTP + WebSocket endpoints)
|
| 254 |
-
└── Dockerfile # Container image definition
|
| 255 |
-
```
|
|
|
|
| 1 |
+
---
|
| 2 |
+
title: Crisisworldcortex Environment Server
|
| 3 |
+
emoji: 🌟
|
| 4 |
+
colorFrom: yellow
|
| 5 |
+
colorTo: pink
|
| 6 |
+
sdk: docker
|
| 7 |
+
pinned: false
|
| 8 |
+
app_port: 8000
|
| 9 |
+
base_path: /web
|
| 10 |
+
tags:
|
| 11 |
+
- openenv
|
| 12 |
+
---
|
| 13 |
+
|
| 14 |
+
# Crisisworldcortex Environment
|
| 15 |
+
|
| 16 |
+
A simple test environment that echoes back messages. Perfect for testing the env APIs as well as demonstrating environment usage patterns.
|
| 17 |
+
|
| 18 |
+
## Quick Start
|
| 19 |
+
|
| 20 |
+
The simplest way to use the Crisisworldcortex environment is through the `CrisisworldcortexEnv` class:
|
| 21 |
+
|
| 22 |
+
```python
|
| 23 |
+
from CrisisWorldCortex import CrisisworldcortexAction, CrisisworldcortexEnv
|
| 24 |
+
|
| 25 |
+
try:
|
| 26 |
+
# Create environment from Docker image
|
| 27 |
+
CrisisWorldCortexenv = CrisisworldcortexEnv.from_docker_image("CrisisWorldCortex-env:latest")
|
| 28 |
+
|
| 29 |
+
# Reset
|
| 30 |
+
result = CrisisWorldCortexenv.reset()
|
| 31 |
+
print(f"Reset: {result.observation.echoed_message}")
|
| 32 |
+
|
| 33 |
+
# Send multiple messages
|
| 34 |
+
messages = ["Hello, World!", "Testing echo", "Final message"]
|
| 35 |
+
|
| 36 |
+
for msg in messages:
|
| 37 |
+
result = CrisisWorldCortexenv.step(CrisisworldcortexAction(message=msg))
|
| 38 |
+
print(f"Sent: '{msg}'")
|
| 39 |
+
print(f" → Echoed: '{result.observation.echoed_message}'")
|
| 40 |
+
print(f" → Length: {result.observation.message_length}")
|
| 41 |
+
print(f" → Reward: {result.reward}")
|
| 42 |
+
|
| 43 |
+
finally:
|
| 44 |
+
# Always clean up
|
| 45 |
+
CrisisWorldCortexenv.close()
|
| 46 |
+
```
|
| 47 |
+
|
| 48 |
+
That's it! The `CrisisworldcortexEnv.from_docker_image()` method handles:
|
| 49 |
+
- Starting the Docker container
|
| 50 |
+
- Waiting for the server to be ready
|
| 51 |
+
- Connecting to the environment
|
| 52 |
+
- Container cleanup when you call `close()`
|
| 53 |
+
|
| 54 |
+
## Building the Docker Image
|
| 55 |
+
|
| 56 |
+
Before using the environment, you need to build the Docker image:
|
| 57 |
+
|
| 58 |
+
```bash
|
| 59 |
+
# From project root
|
| 60 |
+
docker build -t CrisisWorldCortex-env:latest -f server/Dockerfile .
|
| 61 |
+
```
|
| 62 |
+
|
| 63 |
+
## Deploying to Hugging Face Spaces
|
| 64 |
+
|
| 65 |
+
You can easily deploy your OpenEnv environment to Hugging Face Spaces using the `openenv push` command:
|
| 66 |
+
|
| 67 |
+
```bash
|
| 68 |
+
# From the environment directory (where openenv.yaml is located)
|
| 69 |
+
openenv push
|
| 70 |
+
|
| 71 |
+
# Or specify options
|
| 72 |
+
openenv push --namespace my-org --private
|
| 73 |
+
```
|
| 74 |
+
|
| 75 |
+
The `openenv push` command will:
|
| 76 |
+
1. Validate that the directory is an OpenEnv environment (checks for `openenv.yaml`)
|
| 77 |
+
2. Prepare a custom build for Hugging Face Docker space (enables web interface)
|
| 78 |
+
3. Upload to Hugging Face (ensuring you're logged in)
|
| 79 |
+
|
| 80 |
+
### Prerequisites
|
| 81 |
+
|
| 82 |
+
- Authenticate with Hugging Face: The command will prompt for login if not already authenticated
|
| 83 |
+
|
| 84 |
+
### Options
|
| 85 |
+
|
| 86 |
+
- `--directory`, `-d`: Directory containing the OpenEnv environment (defaults to current directory)
|
| 87 |
+
- `--repo-id`, `-r`: Repository ID in format 'username/repo-name' (defaults to 'username/env-name' from openenv.yaml)
|
| 88 |
+
- `--base-image`, `-b`: Base Docker image to use (overrides Dockerfile FROM)
|
| 89 |
+
- `--private`: Deploy the space as private (default: public)
|
| 90 |
+
|
| 91 |
+
### Examples
|
| 92 |
+
|
| 93 |
+
```bash
|
| 94 |
+
# Push to your personal namespace (defaults to username/env-name from openenv.yaml)
|
| 95 |
+
openenv push
|
| 96 |
+
|
| 97 |
+
# Push to a specific repository
|
| 98 |
+
openenv push --repo-id my-org/my-env
|
| 99 |
+
|
| 100 |
+
# Push with a custom base image
|
| 101 |
+
openenv push --base-image ghcr.io/meta-pytorch/openenv-base:latest
|
| 102 |
+
|
| 103 |
+
# Push as a private space
|
| 104 |
+
openenv push --private
|
| 105 |
+
|
| 106 |
+
# Combine options
|
| 107 |
+
openenv push --repo-id my-org/my-env --base-image custom-base:latest --private
|
| 108 |
+
```
|
| 109 |
+
|
| 110 |
+
After deployment, your space will be available at:
|
| 111 |
+
`https://huggingface.co/spaces/<repo-id>`
|
| 112 |
+
|
| 113 |
+
The deployed space includes:
|
| 114 |
+
- **Web Interface** at `/web` - Interactive UI for exploring the environment
|
| 115 |
+
- **API Documentation** at `/docs` - Full OpenAPI/Swagger interface
|
| 116 |
+
- **Health Check** at `/health` - Container health monitoring
|
| 117 |
+
- **WebSocket** at `/ws` - Persistent session endpoint for low-latency interactions
|
| 118 |
+
|
| 119 |
+
## Environment Details
|
| 120 |
+
|
| 121 |
+
### Action
|
| 122 |
+
**CrisisworldcortexAction**: Contains a single field
|
| 123 |
+
- `message` (str) - The message to echo back
|
| 124 |
+
|
| 125 |
+
### Observation
|
| 126 |
+
**CrisisworldcortexObservation**: Contains the echo response and metadata
|
| 127 |
+
- `echoed_message` (str) - The message echoed back
|
| 128 |
+
- `message_length` (int) - Length of the message
|
| 129 |
+
- `reward` (float) - Reward based on message length (length × 0.1)
|
| 130 |
+
- `done` (bool) - Always False for echo environment
|
| 131 |
+
- `metadata` (dict) - Additional info like step count
|
| 132 |
+
|
| 133 |
+
### Reward
|
| 134 |
+
The reward is calculated as: `message_length × 0.1`
|
| 135 |
+
- "Hi" → reward: 0.2
|
| 136 |
+
- "Hello, World!" → reward: 1.3
|
| 137 |
+
- Empty message → reward: 0.0
|
| 138 |
+
|
| 139 |
+
## Advanced Usage
|
| 140 |
+
|
| 141 |
+
### Connecting to an Existing Server
|
| 142 |
+
|
| 143 |
+
If you already have a Crisisworldcortex environment server running, you can connect directly:
|
| 144 |
+
|
| 145 |
+
```python
|
| 146 |
+
from CrisisWorldCortex import CrisisworldcortexEnv
|
| 147 |
+
|
| 148 |
+
# Connect to existing server
|
| 149 |
+
CrisisWorldCortexenv = CrisisworldcortexEnv(base_url="<ENV_HTTP_URL_HERE>")
|
| 150 |
+
|
| 151 |
+
# Use as normal
|
| 152 |
+
result = CrisisWorldCortexenv.reset()
|
| 153 |
+
result = CrisisWorldCortexenv.step(CrisisworldcortexAction(message="Hello!"))
|
| 154 |
+
```
|
| 155 |
+
|
| 156 |
+
Note: When connecting to an existing server, `CrisisWorldCortexenv.close()` will NOT stop the server.
|
| 157 |
+
|
| 158 |
+
### Using the Context Manager
|
| 159 |
+
|
| 160 |
+
The client supports context manager usage for automatic connection management:
|
| 161 |
+
|
| 162 |
+
```python
|
| 163 |
+
from CrisisWorldCortex import CrisisworldcortexAction, CrisisworldcortexEnv
|
| 164 |
+
|
| 165 |
+
# Connect with context manager (auto-connects and closes)
|
| 166 |
+
with CrisisworldcortexEnv(base_url="http://localhost:8000") as env:
|
| 167 |
+
result = env.reset()
|
| 168 |
+
print(f"Reset: {result.observation.echoed_message}")
|
| 169 |
+
# Multiple steps with low latency
|
| 170 |
+
for msg in ["Hello", "World", "!"]:
|
| 171 |
+
result = env.step(CrisisworldcortexAction(message=msg))
|
| 172 |
+
print(f"Echoed: {result.observation.echoed_message}")
|
| 173 |
+
```
|
| 174 |
+
|
| 175 |
+
The client uses WebSocket connections for:
|
| 176 |
+
- **Lower latency**: No HTTP connection overhead per request
|
| 177 |
+
- **Persistent session**: Server maintains your environment state
|
| 178 |
+
- **Efficient for episodes**: Better for many sequential steps
|
| 179 |
+
|
| 180 |
+
### Concurrent WebSocket Sessions
|
| 181 |
+
|
| 182 |
+
The server supports multiple concurrent WebSocket connections. To enable this,
|
| 183 |
+
modify `server/app.py` to use factory mode:
|
| 184 |
+
|
| 185 |
+
```python
|
| 186 |
+
# In server/app.py - use factory mode for concurrent sessions
|
| 187 |
+
app = create_app(
|
| 188 |
+
CrisisworldcortexEnvironment, # Pass class, not instance
|
| 189 |
+
CrisisworldcortexAction,
|
| 190 |
+
CrisisworldcortexObservation,
|
| 191 |
+
max_concurrent_envs=4, # Allow 4 concurrent sessions
|
| 192 |
+
)
|
| 193 |
+
```
|
| 194 |
+
|
| 195 |
+
Then multiple clients can connect simultaneously:
|
| 196 |
+
|
| 197 |
+
```python
|
| 198 |
+
from CrisisWorldCortex import CrisisworldcortexAction, CrisisworldcortexEnv
|
| 199 |
+
from concurrent.futures import ThreadPoolExecutor
|
| 200 |
+
|
| 201 |
+
def run_episode(client_id: int):
|
| 202 |
+
with CrisisworldcortexEnv(base_url="http://localhost:8000") as env:
|
| 203 |
+
result = env.reset()
|
| 204 |
+
for i in range(10):
|
| 205 |
+
result = env.step(CrisisworldcortexAction(message=f"Client {client_id}, step {i}"))
|
| 206 |
+
return client_id, result.observation.message_length
|
| 207 |
+
|
| 208 |
+
# Run 4 episodes concurrently
|
| 209 |
+
with ThreadPoolExecutor(max_workers=4) as executor:
|
| 210 |
+
results = list(executor.map(run_episode, range(4)))
|
| 211 |
+
```
|
| 212 |
+
|
| 213 |
+
## Development & Testing
|
| 214 |
+
|
| 215 |
+
### Direct Environment Testing
|
| 216 |
+
|
| 217 |
+
Test the environment logic directly without starting the HTTP server:
|
| 218 |
+
|
| 219 |
+
```bash
|
| 220 |
+
# From the server directory
|
| 221 |
+
python3 server/CrisisWorldCortex_environment.py
|
| 222 |
+
```
|
| 223 |
+
|
| 224 |
+
This verifies that:
|
| 225 |
+
- Environment resets correctly
|
| 226 |
+
- Step executes actions properly
|
| 227 |
+
- State tracking works
|
| 228 |
+
- Rewards are calculated correctly
|
| 229 |
+
|
| 230 |
+
### Running Locally
|
| 231 |
+
|
| 232 |
+
Run the server locally for development:
|
| 233 |
+
|
| 234 |
+
```bash
|
| 235 |
+
uvicorn server.app:app --reload
|
| 236 |
+
```
|
| 237 |
+
|
| 238 |
+
## Project Structure
|
| 239 |
+
|
| 240 |
+
```
|
| 241 |
+
CrisisWorldCortex/
|
| 242 |
+
├── .dockerignore # Docker build exclusions
|
| 243 |
+
├── __init__.py # Module exports
|
| 244 |
+
├── README.md # This file
|
| 245 |
+
├── openenv.yaml # OpenEnv manifest
|
| 246 |
+
├── pyproject.toml # Project metadata and dependencies
|
| 247 |
+
├── uv.lock # Locked dependencies (generated)
|
| 248 |
+
├── client.py # CrisisworldcortexEnv client
|
| 249 |
+
├── models.py # Action and Observation models
|
| 250 |
+
└── server/
|
| 251 |
+
├── __init__.py # Server module exports
|
| 252 |
+
├── CrisisWorldCortex_environment.py # Core environment logic
|
| 253 |
+
├── app.py # FastAPI application (HTTP + WebSocket endpoints)
|
| 254 |
+
└── Dockerfile # Container image definition
|
| 255 |
+
```
|
__init__.py
CHANGED
|
@@ -1,16 +1,16 @@
|
|
| 1 |
-
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
| 2 |
-
# All rights reserved.
|
| 3 |
-
#
|
| 4 |
-
# This source code is licensed under the BSD-style license found in the
|
| 5 |
-
# LICENSE file in the root directory of this source tree.
|
| 6 |
-
|
| 7 |
-
"""Crisisworldcortex Environment."""
|
| 8 |
-
|
| 9 |
-
from .client import CrisisworldcortexEnv
|
| 10 |
-
from .models import CrisisworldcortexAction, CrisisworldcortexObservation
|
| 11 |
-
|
| 12 |
-
__all__ = [
|
| 13 |
-
"CrisisworldcortexAction",
|
| 14 |
-
"CrisisworldcortexObservation",
|
| 15 |
-
"CrisisworldcortexEnv",
|
| 16 |
-
]
|
|
|
|
| 1 |
+
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
| 2 |
+
# All rights reserved.
|
| 3 |
+
#
|
| 4 |
+
# This source code is licensed under the BSD-style license found in the
|
| 5 |
+
# LICENSE file in the root directory of this source tree.
|
| 6 |
+
|
| 7 |
+
"""Crisisworldcortex Environment."""
|
| 8 |
+
|
| 9 |
+
from .client import CrisisworldcortexEnv
|
| 10 |
+
from .models import CrisisworldcortexAction, CrisisworldcortexObservation
|
| 11 |
+
|
| 12 |
+
__all__ = [
|
| 13 |
+
"CrisisworldcortexAction",
|
| 14 |
+
"CrisisworldcortexObservation",
|
| 15 |
+
"CrisisworldcortexEnv",
|
| 16 |
+
]
|
cortex/brains/__init__.py
CHANGED
|
@@ -0,0 +1,26 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Cortex brains (Session 11+).
|
| 2 |
+
|
| 3 |
+
Public surface:
|
| 4 |
+
- Brain: per-brain class wiring Perception + Lens + 3 LLM subagents + Brain Executive.
|
| 5 |
+
- EpiBrain, LogisticsBrain, GovernanceBrain: factory functions.
|
| 6 |
+
- aggregate_brain_outputs: Brain Executive aggregation function.
|
| 7 |
+
|
| 8 |
+
Each Brain holds its own LLMClient instance; the orchestration layer
|
| 9 |
+
(Council Executive in Session 12, Workstream B trainers) constructs one
|
| 10 |
+
Brain per brain id, optionally with different LLMClients pointing to
|
| 11 |
+
different models. NO module-level state.
|
| 12 |
+
"""
|
| 13 |
+
|
| 14 |
+
from ._base import Brain
|
| 15 |
+
from ._executive import aggregate_brain_outputs
|
| 16 |
+
from .epidemiology import EpiBrain
|
| 17 |
+
from .governance import GovernanceBrain
|
| 18 |
+
from .logistics import LogisticsBrain
|
| 19 |
+
|
| 20 |
+
__all__ = [
|
| 21 |
+
"Brain",
|
| 22 |
+
"EpiBrain",
|
| 23 |
+
"GovernanceBrain",
|
| 24 |
+
"LogisticsBrain",
|
| 25 |
+
"aggregate_brain_outputs",
|
| 26 |
+
]
|
cortex/brains/_base.py
ADDED
|
@@ -0,0 +1,201 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Brain class - composes Perception + Lens + 3 Subagents + Brain Executive.
|
| 2 |
+
|
| 3 |
+
Per cortex/CLAUDE.md: each brain has a deterministic Python Perception
|
| 4 |
+
+ Lens, three LLM subagents (router-callable), and a deterministic
|
| 5 |
+
Python Brain Executive. The Brain class wires these together.
|
| 6 |
+
|
| 7 |
+
Multi-model deployment: each Brain holds a SINGLE LLMClient instance
|
| 8 |
+
passed at construction. Different brains can use different models by
|
| 9 |
+
constructing each Brain with a different LLMClient (e.g., Qwen for epi,
|
| 10 |
+
Llama for logistics). NO module-level state, NO shared singletons.
|
| 11 |
+
"""
|
| 12 |
+
|
| 13 |
+
from __future__ import annotations
|
| 14 |
+
|
| 15 |
+
from typing import List, Literal
|
| 16 |
+
|
| 17 |
+
from cortex.lenses import lens_for
|
| 18 |
+
from cortex.schemas import (
|
| 19 |
+
BeliefState,
|
| 20 |
+
BrainLensedObservation,
|
| 21 |
+
BrainRecommendation,
|
| 22 |
+
CandidatePlan,
|
| 23 |
+
CriticReport,
|
| 24 |
+
PerceptionReport,
|
| 25 |
+
SubagentInput,
|
| 26 |
+
)
|
| 27 |
+
from cortex.subagents import (
|
| 28 |
+
CriticSubagent,
|
| 29 |
+
PlannerSubagent,
|
| 30 |
+
WorldModelerSubagent,
|
| 31 |
+
perception_for,
|
| 32 |
+
)
|
| 33 |
+
from cortex.subagents._base import _LLMClientLike
|
| 34 |
+
from CrisisWorldCortex.models import CrisisworldcortexObservation
|
| 35 |
+
|
| 36 |
+
from ._executive import aggregate_brain_outputs
|
| 37 |
+
|
| 38 |
+
_BrainId = Literal["epidemiology", "logistics", "governance"]
|
| 39 |
+
|
| 40 |
+
|
| 41 |
+
class Brain:
|
| 42 |
+
"""Per-brain pipeline holder.
|
| 43 |
+
|
| 44 |
+
Each Brain instance owns its own LLMClient. The orchestration layer
|
| 45 |
+
(Session 12 Council, Workstream B trainers) instantiates one Brain
|
| 46 |
+
per brain id, optionally with different LLMClients pointing to
|
| 47 |
+
different models. The Brain class itself has NO module-level state
|
| 48 |
+
and NO forced singleton.
|
| 49 |
+
|
| 50 |
+
Args:
|
| 51 |
+
brain_id: One of "epidemiology", "logistics", "governance".
|
| 52 |
+
llm_client: This brain's LLM client. Subagents are constructed
|
| 53 |
+
with the SAME client so token billing aggregates correctly.
|
| 54 |
+
wm: WorldModeler subagent.
|
| 55 |
+
planner: Planner subagent.
|
| 56 |
+
critic: Critic subagent.
|
| 57 |
+
"""
|
| 58 |
+
|
| 59 |
+
def __init__(
|
| 60 |
+
self,
|
| 61 |
+
brain_id: _BrainId,
|
| 62 |
+
llm_client: _LLMClientLike,
|
| 63 |
+
wm: WorldModelerSubagent,
|
| 64 |
+
planner: PlannerSubagent,
|
| 65 |
+
critic: CriticSubagent,
|
| 66 |
+
) -> None:
|
| 67 |
+
self.brain_id = brain_id
|
| 68 |
+
self.llm_client = llm_client
|
| 69 |
+
self.wm = wm
|
| 70 |
+
self.planner = planner
|
| 71 |
+
self.critic = critic
|
| 72 |
+
|
| 73 |
+
# ------------------------------------------------------------------
|
| 74 |
+
# Deterministic Python pieces (no LLM)
|
| 75 |
+
# ------------------------------------------------------------------
|
| 76 |
+
|
| 77 |
+
def compute_perception(self, obs: CrisisworldcortexObservation) -> PerceptionReport:
|
| 78 |
+
"""Run this brain's Perception. Pure Python; no LLM."""
|
| 79 |
+
return perception_for(self.brain_id, obs)
|
| 80 |
+
|
| 81 |
+
def compute_lens(
|
| 82 |
+
self, obs: CrisisworldcortexObservation, last_reward: float
|
| 83 |
+
) -> BrainLensedObservation:
|
| 84 |
+
"""Run this brain's Lens. Pure Python; no LLM."""
|
| 85 |
+
return lens_for(self.brain_id, obs, last_reward)
|
| 86 |
+
|
| 87 |
+
def aggregate(
|
| 88 |
+
self,
|
| 89 |
+
perception: PerceptionReport,
|
| 90 |
+
beliefs: List[BeliefState],
|
| 91 |
+
plans: List[CandidatePlan],
|
| 92 |
+
critics: List[CriticReport],
|
| 93 |
+
tokens_used: int = 0,
|
| 94 |
+
) -> BrainRecommendation:
|
| 95 |
+
"""Run this brain's Brain Executive. Pure Python; no LLM."""
|
| 96 |
+
return aggregate_brain_outputs(
|
| 97 |
+
brain_id=self.brain_id,
|
| 98 |
+
perception=perception,
|
| 99 |
+
beliefs=beliefs,
|
| 100 |
+
plans=plans,
|
| 101 |
+
critics=critics,
|
| 102 |
+
tokens_used=tokens_used,
|
| 103 |
+
)
|
| 104 |
+
|
| 105 |
+
# ------------------------------------------------------------------
|
| 106 |
+
# High-level convenience: round-1 single tick
|
| 107 |
+
# ------------------------------------------------------------------
|
| 108 |
+
|
| 109 |
+
def run_tick(
|
| 110 |
+
self,
|
| 111 |
+
obs: CrisisworldcortexObservation,
|
| 112 |
+
last_reward: float,
|
| 113 |
+
tick: int,
|
| 114 |
+
round_: int = 1,
|
| 115 |
+
) -> BrainRecommendation:
|
| 116 |
+
"""Round-1 single-tick pipeline (Session 11 smoke).
|
| 117 |
+
|
| 118 |
+
Round 2 is orchestrated by the Council Executive (Session 12)
|
| 119 |
+
via the fine-grained methods (compute_perception, compute_lens,
|
| 120 |
+
wm.run / planner.run / critic.run, aggregate). Calling this
|
| 121 |
+
convenience method with ``round_!=1`` raises NotImplementedError
|
| 122 |
+
to prevent accidental misuse before the Council exists.
|
| 123 |
+
"""
|
| 124 |
+
if round_ != 1:
|
| 125 |
+
raise NotImplementedError(
|
| 126 |
+
f"Round {round_} orchestration is the Council Executive's "
|
| 127 |
+
f"responsibility (Session 12). Use Brain.compute_perception/"
|
| 128 |
+
f"compute_lens + WorldModelerSubagent.run/PlannerSubagent.run/"
|
| 129 |
+
f"CriticSubagent.run + Brain.aggregate directly."
|
| 130 |
+
)
|
| 131 |
+
|
| 132 |
+
perception = self.compute_perception(obs)
|
| 133 |
+
# Lens is computed for completeness; Session 11 doesn't yet plumb
|
| 134 |
+
# it into SubagentInput (M-FR-4 step indices fixed). Session 12
|
| 135 |
+
# Council will extend the SubagentInput contract to carry lens
|
| 136 |
+
# output if subagents need it.
|
| 137 |
+
_ = self.compute_lens(obs, last_reward)
|
| 138 |
+
|
| 139 |
+
# WorldModeler (step_idx=0)
|
| 140 |
+
wm_input = SubagentInput(
|
| 141 |
+
brain=self.brain_id,
|
| 142 |
+
role="world_modeler",
|
| 143 |
+
tick=tick,
|
| 144 |
+
round=round_,
|
| 145 |
+
perception=perception,
|
| 146 |
+
prior_belief=None,
|
| 147 |
+
prior_plans=[],
|
| 148 |
+
target_plan_id=None,
|
| 149 |
+
last_reward=last_reward,
|
| 150 |
+
recent_action_log_excerpt=list(obs.recent_action_log),
|
| 151 |
+
)
|
| 152 |
+
belief = self.wm.run(wm_input, step_idx=0)
|
| 153 |
+
|
| 154 |
+
# Planner (step_idx=1)
|
| 155 |
+
planner_input = SubagentInput(
|
| 156 |
+
brain=self.brain_id,
|
| 157 |
+
role="planner",
|
| 158 |
+
tick=tick,
|
| 159 |
+
round=round_,
|
| 160 |
+
perception=perception,
|
| 161 |
+
prior_belief=belief,
|
| 162 |
+
prior_plans=[],
|
| 163 |
+
target_plan_id=None,
|
| 164 |
+
last_reward=last_reward,
|
| 165 |
+
recent_action_log_excerpt=list(obs.recent_action_log),
|
| 166 |
+
)
|
| 167 |
+
plan = self.planner.run(planner_input, step_idx=1)
|
| 168 |
+
|
| 169 |
+
# Critic (step_idx=2)
|
| 170 |
+
critic_input = SubagentInput(
|
| 171 |
+
brain=self.brain_id,
|
| 172 |
+
role="critic",
|
| 173 |
+
tick=tick,
|
| 174 |
+
round=round_,
|
| 175 |
+
perception=perception,
|
| 176 |
+
prior_belief=belief,
|
| 177 |
+
prior_plans=[plan],
|
| 178 |
+
target_plan_id="plan-0",
|
| 179 |
+
last_reward=last_reward,
|
| 180 |
+
recent_action_log_excerpt=list(obs.recent_action_log),
|
| 181 |
+
)
|
| 182 |
+
critic = self.critic.run(critic_input, step_idx=2)
|
| 183 |
+
|
| 184 |
+
# Tally tokens billed to this brain's caller_ids.
|
| 185 |
+
caller_id_base = f"cortex:{self.brain_id}"
|
| 186 |
+
tokens_used = sum(
|
| 187 |
+
self.llm_client.tokens_used_for(f"{caller_id_base}:{role}:t{tick}:r{round_}:s{idx}")
|
| 188 |
+
for role, idx in (
|
| 189 |
+
("world_modeler", 0),
|
| 190 |
+
("planner", 1),
|
| 191 |
+
("critic", 2),
|
| 192 |
+
)
|
| 193 |
+
)
|
| 194 |
+
|
| 195 |
+
return self.aggregate(
|
| 196 |
+
perception=perception,
|
| 197 |
+
beliefs=[belief],
|
| 198 |
+
plans=[plan],
|
| 199 |
+
critics=[critic],
|
| 200 |
+
tokens_used=tokens_used,
|
| 201 |
+
)
|
cortex/brains/_executive.py
ADDED
|
@@ -0,0 +1,144 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Brain Executive - deterministic Python aggregation.
|
| 2 |
+
|
| 3 |
+
Per Phase A docs/CORTEX_ARCHITECTURE.md Decisions 15-21 + M-FR-3 partial
|
| 4 |
+
evidence union (perception + beliefs only; CandidatePlan and CriticReport
|
| 5 |
+
schemas have no evidence fields).
|
| 6 |
+
|
| 7 |
+
Brain Executive runs ONCE per brain at round end. NOT router-callable
|
| 8 |
+
per cortex/CLAUDE.md.
|
| 9 |
+
"""
|
| 10 |
+
|
| 11 |
+
from __future__ import annotations
|
| 12 |
+
|
| 13 |
+
from typing import List
|
| 14 |
+
|
| 15 |
+
from cortex.schemas import (
|
| 16 |
+
BeliefState,
|
| 17 |
+
BrainRecommendation,
|
| 18 |
+
CandidatePlan,
|
| 19 |
+
CriticReport,
|
| 20 |
+
EvidenceCitation,
|
| 21 |
+
PerceptionReport,
|
| 22 |
+
)
|
| 23 |
+
from CrisisWorldCortex.models import NoOp
|
| 24 |
+
|
| 25 |
+
_REASONING_SUMMARY_MAX_CHARS = 400 # matches BrainRecommendation.reasoning_summary cap
|
| 26 |
+
_FALSIFIERS_TO_JOIN = 3
|
| 27 |
+
_FALSIFIER_FALLBACK = "(no falsifier provided)"
|
| 28 |
+
_EMPTY_REASONING = "(empty: no subagent produced a parseable plan)"
|
| 29 |
+
|
| 30 |
+
|
| 31 |
+
def aggregate_brain_outputs(
|
| 32 |
+
brain_id: str,
|
| 33 |
+
perception: PerceptionReport,
|
| 34 |
+
beliefs: List[BeliefState],
|
| 35 |
+
plans: List[CandidatePlan],
|
| 36 |
+
critics: List[CriticReport],
|
| 37 |
+
tokens_used: int = 0,
|
| 38 |
+
) -> BrainRecommendation:
|
| 39 |
+
"""Aggregate one brain's per-round subagent outputs into a recommendation.
|
| 40 |
+
|
| 41 |
+
Decisions:
|
| 42 |
+
D15: argmax over expected_value * confidence.
|
| 43 |
+
D16: top_confidence = chosen.confidence * (1 - chosen_belief.uncertainty).
|
| 44 |
+
D17: minority_actions = all expected_outer_actions except chosen.
|
| 45 |
+
D19: reasoning_summary = chosen.action_sketch[:400].
|
| 46 |
+
D20 (M-FR-3): evidence = perception.evidence + flat-union of beliefs[*].evidence.
|
| 47 |
+
CandidatePlan/CriticReport carry no evidence fields.
|
| 48 |
+
D21: brain_id is lowercase per Pydantic Literal in BrainRecommendation.
|
| 49 |
+
|
| 50 |
+
Empty fallback (M-FR-7): no plans, or chosen plan has confidence==0
|
| 51 |
+
-> top_action=NoOp, top_confidence=0, uncertainty=1.0,
|
| 52 |
+
reasoning_summary=_EMPTY_REASONING.
|
| 53 |
+
|
| 54 |
+
Args:
|
| 55 |
+
brain_id: lowercase brain id ("epidemiology" / "logistics" / "governance").
|
| 56 |
+
perception: This brain's PerceptionReport for the tick.
|
| 57 |
+
beliefs: Per-round BeliefStates. Index aligned with ``plans``.
|
| 58 |
+
plans: Per-round CandidatePlans.
|
| 59 |
+
critics: Per-round CriticReports (currently unused in aggregation but
|
| 60 |
+
kept on the signature so the trajectory log captures the full
|
| 61 |
+
chain).
|
| 62 |
+
tokens_used: Total tokens billed across this brain's subagents.
|
| 63 |
+
"""
|
| 64 |
+
if not plans:
|
| 65 |
+
return _empty_recommendation(brain_id, perception, beliefs, tokens_used)
|
| 66 |
+
|
| 67 |
+
# D15: argmax over expected_value * confidence
|
| 68 |
+
chosen_idx = max(
|
| 69 |
+
range(len(plans)),
|
| 70 |
+
key=lambda i: plans[i].expected_value * plans[i].confidence,
|
| 71 |
+
)
|
| 72 |
+
chosen_plan = plans[chosen_idx]
|
| 73 |
+
|
| 74 |
+
if chosen_plan.confidence == 0.0:
|
| 75 |
+
# All plans are empty fallbacks (or the only plan is empty).
|
| 76 |
+
# Brain Executive treats this as no-signal.
|
| 77 |
+
return _empty_recommendation(brain_id, perception, beliefs, tokens_used)
|
| 78 |
+
|
| 79 |
+
# D16: top_confidence = chosen.confidence * (1 - belief.uncertainty)
|
| 80 |
+
if chosen_idx < len(beliefs):
|
| 81 |
+
chosen_belief = beliefs[chosen_idx]
|
| 82 |
+
uncertainty = chosen_belief.uncertainty
|
| 83 |
+
else:
|
| 84 |
+
# Defensive: parallel arrays should match. If not, treat as max uncertainty.
|
| 85 |
+
uncertainty = 1.0
|
| 86 |
+
top_confidence = chosen_plan.confidence * (1.0 - uncertainty)
|
| 87 |
+
|
| 88 |
+
# D17: minority_actions = all plans except chosen
|
| 89 |
+
minority_actions = [
|
| 90 |
+
plans[i].expected_outer_action for i in range(len(plans)) if i != chosen_idx
|
| 91 |
+
]
|
| 92 |
+
|
| 93 |
+
# D19: reasoning_summary
|
| 94 |
+
reasoning_summary = chosen_plan.action_sketch[:_REASONING_SUMMARY_MAX_CHARS]
|
| 95 |
+
|
| 96 |
+
# D20 (M-FR-3): evidence union from perception + beliefs only.
|
| 97 |
+
# CandidatePlan and CriticReport schemas (Session 9) have no evidence fields;
|
| 98 |
+
# the perception+beliefs union captures the actionable evidence chain since
|
| 99 |
+
# plans/critics derive from beliefs.
|
| 100 |
+
evidence: List[EvidenceCitation] = list(perception.evidence)
|
| 101 |
+
for b in beliefs:
|
| 102 |
+
evidence.extend(b.evidence)
|
| 103 |
+
|
| 104 |
+
# falsifier (M-FR-6): join up to 3 falsifiers; fallback if empty.
|
| 105 |
+
if chosen_plan.falsifiers:
|
| 106 |
+
falsifier = "; ".join(chosen_plan.falsifiers[:_FALSIFIERS_TO_JOIN])
|
| 107 |
+
else:
|
| 108 |
+
falsifier = _FALSIFIER_FALLBACK
|
| 109 |
+
|
| 110 |
+
return BrainRecommendation(
|
| 111 |
+
brain=brain_id,
|
| 112 |
+
top_action=chosen_plan.expected_outer_action,
|
| 113 |
+
top_confidence=top_confidence,
|
| 114 |
+
minority_actions=minority_actions,
|
| 115 |
+
reasoning_summary=reasoning_summary,
|
| 116 |
+
evidence=evidence,
|
| 117 |
+
falsifier=falsifier,
|
| 118 |
+
uncertainty=uncertainty,
|
| 119 |
+
tokens_used=tokens_used,
|
| 120 |
+
)
|
| 121 |
+
|
| 122 |
+
|
| 123 |
+
def _empty_recommendation(
|
| 124 |
+
brain_id: str,
|
| 125 |
+
perception: PerceptionReport,
|
| 126 |
+
beliefs: List[BeliefState],
|
| 127 |
+
tokens_used: int,
|
| 128 |
+
) -> BrainRecommendation:
|
| 129 |
+
"""M-FR-7 empty fallback: NoOp + confidence=0 + uncertainty=1."""
|
| 130 |
+
evidence: List[EvidenceCitation] = list(perception.evidence)
|
| 131 |
+
for b in beliefs:
|
| 132 |
+
evidence.extend(b.evidence)
|
| 133 |
+
|
| 134 |
+
return BrainRecommendation(
|
| 135 |
+
brain=brain_id,
|
| 136 |
+
top_action=NoOp(),
|
| 137 |
+
top_confidence=0.0,
|
| 138 |
+
minority_actions=[],
|
| 139 |
+
reasoning_summary=_EMPTY_REASONING,
|
| 140 |
+
evidence=evidence,
|
| 141 |
+
falsifier=_FALSIFIER_FALLBACK,
|
| 142 |
+
uncertainty=1.0,
|
| 143 |
+
tokens_used=tokens_used,
|
| 144 |
+
)
|
cortex/brains/epidemiology.py
ADDED
|
@@ -0,0 +1,25 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Epidemiology brain factory."""
|
| 2 |
+
|
| 3 |
+
from __future__ import annotations
|
| 4 |
+
|
| 5 |
+
from cortex.subagents import CriticSubagent, PlannerSubagent, WorldModelerSubagent
|
| 6 |
+
from cortex.subagents._base import _LLMClientLike
|
| 7 |
+
|
| 8 |
+
from ._base import Brain
|
| 9 |
+
|
| 10 |
+
|
| 11 |
+
def EpiBrain(llm_client: _LLMClientLike) -> Brain:
|
| 12 |
+
"""Construct an Epidemiology Brain bound to ``llm_client``.
|
| 13 |
+
|
| 14 |
+
Multi-model deployment: pass a different ``llm_client`` per brain
|
| 15 |
+
instance to use different models per brain (e.g., Qwen for epi,
|
| 16 |
+
Llama for logistics). The 3 LLM subagents are constructed with the
|
| 17 |
+
SAME client so token billing aggregates correctly.
|
| 18 |
+
"""
|
| 19 |
+
return Brain(
|
| 20 |
+
brain_id="epidemiology",
|
| 21 |
+
llm_client=llm_client,
|
| 22 |
+
wm=WorldModelerSubagent(llm_client),
|
| 23 |
+
planner=PlannerSubagent(llm_client),
|
| 24 |
+
critic=CriticSubagent(llm_client),
|
| 25 |
+
)
|
cortex/brains/governance.py
ADDED
|
@@ -0,0 +1,22 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Governance brain factory."""
|
| 2 |
+
|
| 3 |
+
from __future__ import annotations
|
| 4 |
+
|
| 5 |
+
from cortex.subagents import CriticSubagent, PlannerSubagent, WorldModelerSubagent
|
| 6 |
+
from cortex.subagents._base import _LLMClientLike
|
| 7 |
+
|
| 8 |
+
from ._base import Brain
|
| 9 |
+
|
| 10 |
+
|
| 11 |
+
def GovernanceBrain(llm_client: _LLMClientLike) -> Brain:
|
| 12 |
+
"""Construct a Governance Brain bound to ``llm_client``.
|
| 13 |
+
|
| 14 |
+
Multi-model deployment: see EpiBrain.
|
| 15 |
+
"""
|
| 16 |
+
return Brain(
|
| 17 |
+
brain_id="governance",
|
| 18 |
+
llm_client=llm_client,
|
| 19 |
+
wm=WorldModelerSubagent(llm_client),
|
| 20 |
+
planner=PlannerSubagent(llm_client),
|
| 21 |
+
critic=CriticSubagent(llm_client),
|
| 22 |
+
)
|
cortex/brains/logistics.py
ADDED
|
@@ -0,0 +1,23 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Logistics brain factory."""
|
| 2 |
+
|
| 3 |
+
from __future__ import annotations
|
| 4 |
+
|
| 5 |
+
from cortex.subagents import CriticSubagent, PlannerSubagent, WorldModelerSubagent
|
| 6 |
+
from cortex.subagents._base import _LLMClientLike
|
| 7 |
+
|
| 8 |
+
from ._base import Brain
|
| 9 |
+
|
| 10 |
+
|
| 11 |
+
def LogisticsBrain(llm_client: _LLMClientLike) -> Brain:
|
| 12 |
+
"""Construct a Logistics Brain bound to ``llm_client``.
|
| 13 |
+
|
| 14 |
+
Multi-model deployment: see EpiBrain. Pass a Llama-bound client
|
| 15 |
+
here while the other brains use Qwen, etc.
|
| 16 |
+
"""
|
| 17 |
+
return Brain(
|
| 18 |
+
brain_id="logistics",
|
| 19 |
+
llm_client=llm_client,
|
| 20 |
+
wm=WorldModelerSubagent(llm_client),
|
| 21 |
+
planner=PlannerSubagent(llm_client),
|
| 22 |
+
critic=CriticSubagent(llm_client),
|
| 23 |
+
)
|
cortex/lenses.py
ADDED
|
@@ -0,0 +1,186 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Brain-specific observation lenses (Session 10).
|
| 2 |
+
|
| 3 |
+
Phase A docs/CORTEX_ARCHITECTURE.md Decisions 9-14 + §2 A1.
|
| 4 |
+
|
| 5 |
+
Lenses are pure-Python: no LLM, no I/O, no state. ``lens_for(brain, obs,
|
| 6 |
+
last_reward)`` dispatches to one of 3 brain-specific helpers and returns
|
| 7 |
+
a ``BrainLensedObservation``. V2 brain ids raise ``KeyError`` per
|
| 8 |
+
Decision 9 (post-review) -- no MVP stub functions.
|
| 9 |
+
|
| 10 |
+
The lens does NOT strip the raw observation (Decision 13); subagents
|
| 11 |
+
may need fields the lens didn't emphasise. ``salient_field_ids`` is a
|
| 12 |
+
salience map alongside ``raw_obs``, not a replacement for it.
|
| 13 |
+
|
| 14 |
+
``transmission_rate_trend`` is fixed at 0.0 in MVP (M-FR-2): the lens
|
| 15 |
+
sees one observation per call. Session 11 plumbs prior-tick obs into
|
| 16 |
+
the lens to enable real trend computation.
|
| 17 |
+
"""
|
| 18 |
+
|
| 19 |
+
from __future__ import annotations
|
| 20 |
+
|
| 21 |
+
from typing import Callable, Dict
|
| 22 |
+
|
| 23 |
+
from cortex.schemas import BrainLensedObservation
|
| 24 |
+
from CrisisWorldCortex.models import CrisisworldcortexObservation
|
| 25 |
+
|
| 26 |
+
_V2_BRAINS = frozenset({"communications", "equity"})
|
| 27 |
+
|
| 28 |
+
|
| 29 |
+
def lens_for(
|
| 30 |
+
brain: str,
|
| 31 |
+
obs: CrisisworldcortexObservation,
|
| 32 |
+
last_reward: float,
|
| 33 |
+
) -> BrainLensedObservation:
|
| 34 |
+
"""Return the per-brain lensed observation.
|
| 35 |
+
|
| 36 |
+
Args:
|
| 37 |
+
brain: One of {"epidemiology", "logistics", "governance"}.
|
| 38 |
+
obs: The current tick's observation.
|
| 39 |
+
last_reward: Previous tick's reward (plumbed from B1's pattern,
|
| 40 |
+
included on the lensed object so all 3 subagents in this
|
| 41 |
+
brain see the same recency signal).
|
| 42 |
+
|
| 43 |
+
Raises:
|
| 44 |
+
KeyError: If ``brain`` is a V2-deferred brain (communications,
|
| 45 |
+
equity) or unknown.
|
| 46 |
+
"""
|
| 47 |
+
helper = _LENS_REGISTRY.get(brain)
|
| 48 |
+
if helper is not None:
|
| 49 |
+
return helper(obs, last_reward)
|
| 50 |
+
if brain in _V2_BRAINS:
|
| 51 |
+
raise KeyError(
|
| 52 |
+
f"V2 brain {brain!r} deferred per Phase A Decision 9; "
|
| 53 |
+
f"no MVP stub lens. See docs/CORTEX_ARCHITECTURE.md."
|
| 54 |
+
)
|
| 55 |
+
raise KeyError(f"unknown brain: {brain!r}")
|
| 56 |
+
|
| 57 |
+
|
| 58 |
+
# ============================================================================
|
| 59 |
+
# Per-brain lens helpers
|
| 60 |
+
# ============================================================================
|
| 61 |
+
|
| 62 |
+
|
| 63 |
+
def _epi_lens(obs: CrisisworldcortexObservation, last_reward: float) -> BrainLensedObservation:
|
| 64 |
+
"""Epidemiology lens (Decision 10, M-FR-4 rename: epi_pressure)."""
|
| 65 |
+
n_regions = max(1, len(obs.regions))
|
| 66 |
+
mean_hospital_load = sum(r.hospital_load for r in obs.regions) / n_regions
|
| 67 |
+
# M-FR-4: pressure scalar correlated with R_eff but not a true R_eff
|
| 68 |
+
# estimate. WorldModeler subagent computes proper R_eff during reasoning.
|
| 69 |
+
epi_pressure = max(0.0, min(3.0, mean_hospital_load * 2.0))
|
| 70 |
+
|
| 71 |
+
max_cases = max((r.reported_cases_d_ago for r in obs.regions), default=0)
|
| 72 |
+
# /1000 normaliser matches the design-doc "~30 cases / 1000 pop" spec
|
| 73 |
+
worst_region_infection = max(0.0, min(1.0, max_cases / 1000.0))
|
| 74 |
+
|
| 75 |
+
return BrainLensedObservation(
|
| 76 |
+
brain="epidemiology",
|
| 77 |
+
raw_obs=obs,
|
| 78 |
+
salient_field_ids=[
|
| 79 |
+
"regions[*].reported_cases_d_ago",
|
| 80 |
+
"regions[*].hospital_load",
|
| 81 |
+
"regions[*].compliance_proxy",
|
| 82 |
+
],
|
| 83 |
+
derived_features={
|
| 84 |
+
"epi_pressure": float(epi_pressure),
|
| 85 |
+
"worst_region_infection": float(worst_region_infection),
|
| 86 |
+
# M-FR-2: needs history; Session 11 plumbs prior-tick obs.
|
| 87 |
+
"transmission_rate_trend": 0.0,
|
| 88 |
+
},
|
| 89 |
+
last_reward=last_reward,
|
| 90 |
+
)
|
| 91 |
+
|
| 92 |
+
|
| 93 |
+
def _logistics_lens(
|
| 94 |
+
obs: CrisisworldcortexObservation, last_reward: float
|
| 95 |
+
) -> BrainLensedObservation:
|
| 96 |
+
"""Logistics lens (Decision 11, M-FR-3 floor 0.5, M-FR-6 flat keys)."""
|
| 97 |
+
res = obs.resources
|
| 98 |
+
total_inventory = float(
|
| 99 |
+
res.test_kits + res.hospital_beds_free + res.mobile_units + res.vaccine_doses
|
| 100 |
+
)
|
| 101 |
+
|
| 102 |
+
hospital_load_max = (
|
| 103 |
+
max((r.hospital_load for r in obs.regions), default=0.0) if obs.regions else 0.0
|
| 104 |
+
)
|
| 105 |
+
|
| 106 |
+
# Per-region feasibility flat keys (D14 + M-FR-6).
|
| 107 |
+
strict_regions = {r.region for r in obs.active_restrictions if r.severity == "strict"}
|
| 108 |
+
feasibility: Dict[str, float] = {}
|
| 109 |
+
for r in obs.regions:
|
| 110 |
+
key = f"deployment_feasibility_{r.region}"
|
| 111 |
+
if total_inventory <= 0.0:
|
| 112 |
+
feasibility[key] = 0.0
|
| 113 |
+
elif r.region in strict_regions:
|
| 114 |
+
# M-FR-3: 0.5 floor when strict restriction is in place but
|
| 115 |
+
# units could still be helicoptered in; Planner does the
|
| 116 |
+
# legal-check.
|
| 117 |
+
feasibility[key] = 0.5
|
| 118 |
+
else:
|
| 119 |
+
feasibility[key] = 1.0
|
| 120 |
+
|
| 121 |
+
derived_features: Dict[str, float] = {
|
| 122 |
+
"total_inventory": total_inventory,
|
| 123 |
+
"hospital_load_max": float(hospital_load_max),
|
| 124 |
+
**feasibility,
|
| 125 |
+
}
|
| 126 |
+
|
| 127 |
+
return BrainLensedObservation(
|
| 128 |
+
brain="logistics",
|
| 129 |
+
raw_obs=obs,
|
| 130 |
+
salient_field_ids=[
|
| 131 |
+
"resources.test_kits",
|
| 132 |
+
"resources.hospital_beds_free",
|
| 133 |
+
"resources.mobile_units",
|
| 134 |
+
"resources.vaccine_doses",
|
| 135 |
+
"regions[*].hospital_load",
|
| 136 |
+
"active_restrictions[*]",
|
| 137 |
+
],
|
| 138 |
+
derived_features=derived_features,
|
| 139 |
+
last_reward=last_reward,
|
| 140 |
+
)
|
| 141 |
+
|
| 142 |
+
|
| 143 |
+
def _governance_lens(
|
| 144 |
+
obs: CrisisworldcortexObservation, last_reward: float
|
| 145 |
+
) -> BrainLensedObservation:
|
| 146 |
+
"""Governance lens (Decision 12)."""
|
| 147 |
+
# escalation_unlocked_strict: any accepted escalate(national) in the log
|
| 148 |
+
escalation_unlocked = 0.0
|
| 149 |
+
for ea in obs.recent_action_log:
|
| 150 |
+
if (
|
| 151 |
+
ea.accepted
|
| 152 |
+
and ea.action.kind == "escalate"
|
| 153 |
+
and getattr(ea.action, "to_authority", None) == "national"
|
| 154 |
+
):
|
| 155 |
+
escalation_unlocked = 1.0
|
| 156 |
+
break
|
| 157 |
+
|
| 158 |
+
return BrainLensedObservation(
|
| 159 |
+
brain="governance",
|
| 160 |
+
raw_obs=obs,
|
| 161 |
+
salient_field_ids=[
|
| 162 |
+
"active_restrictions[*]",
|
| 163 |
+
"legal_constraints[*]",
|
| 164 |
+
"recent_action_log[*]",
|
| 165 |
+
],
|
| 166 |
+
derived_features={
|
| 167 |
+
"escalation_unlocked_strict": escalation_unlocked,
|
| 168 |
+
"legal_constraints_count": float(len(obs.legal_constraints)),
|
| 169 |
+
"restrictions_active_count": float(len(obs.active_restrictions)),
|
| 170 |
+
},
|
| 171 |
+
last_reward=last_reward,
|
| 172 |
+
)
|
| 173 |
+
|
| 174 |
+
|
| 175 |
+
# ============================================================================
|
| 176 |
+
# Registry (defined after helpers so closures resolve cleanly)
|
| 177 |
+
# ============================================================================
|
| 178 |
+
|
| 179 |
+
|
| 180 |
+
_LENS_REGISTRY: Dict[
|
| 181 |
+
str, Callable[[CrisisworldcortexObservation, float], BrainLensedObservation]
|
| 182 |
+
] = {
|
| 183 |
+
"epidemiology": _epi_lens,
|
| 184 |
+
"logistics": _logistics_lens,
|
| 185 |
+
"governance": _governance_lens,
|
| 186 |
+
}
|
cortex/schemas.py
CHANGED
|
@@ -35,7 +35,12 @@ from pydantic import BaseModel, Field
|
|
| 35 |
# OWN internal types (cortex.subagents, cortex.brains, etc.) continue
|
| 36 |
# to use bare-name sibling imports per Phase 1 C1 — only the cross-package
|
| 37 |
# wire boundary is canonicalised.
|
| 38 |
-
from CrisisWorldCortex.models import
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 39 |
|
| 40 |
EpistemicPhase = Literal["Divergence", "Challenge", "Narrowing", "Convergence"]
|
| 41 |
|
|
@@ -137,6 +142,46 @@ SubagentReport = Union[BeliefState, CandidatePlan, CriticReport]
|
|
| 137 |
trajectory buffers that need to carry 'any subagent output' generically."""
|
| 138 |
|
| 139 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 140 |
# ============================================================================
|
| 141 |
# Brain output + Council decision
|
| 142 |
# ============================================================================
|
|
|
|
| 35 |
# OWN internal types (cortex.subagents, cortex.brains, etc.) continue
|
| 36 |
# to use bare-name sibling imports per Phase 1 C1 — only the cross-package
|
| 37 |
# wire boundary is canonicalised.
|
| 38 |
+
from CrisisWorldCortex.models import (
|
| 39 |
+
CrisisworldcortexObservation,
|
| 40 |
+
ExecutedAction,
|
| 41 |
+
OuterActionPayload,
|
| 42 |
+
RegionId,
|
| 43 |
+
)
|
| 44 |
|
| 45 |
EpistemicPhase = Literal["Divergence", "Challenge", "Narrowing", "Convergence"]
|
| 46 |
|
|
|
|
| 142 |
trajectory buffers that need to carry 'any subagent output' generically."""
|
| 143 |
|
| 144 |
|
| 145 |
+
class SubagentInput(BaseModel):
|
| 146 |
+
"""Typed input handed to one of the 3 LLM subagents per call.
|
| 147 |
+
|
| 148 |
+
Per Phase A §2 A2: each subagent call receives a fully-typed input
|
| 149 |
+
so prompts are deterministic and testable. ``prior_belief`` is
|
| 150 |
+
``None`` on round 1 (nothing to revise yet); on round 2 it carries
|
| 151 |
+
the previous round's BeliefState (or an empty BeliefState if round 1
|
| 152 |
+
failed, per Phase A Decision 62). ``prior_plans`` is empty for
|
| 153 |
+
WorldModeler / Planner; populated for Critic so it can attack a
|
| 154 |
+
specific plan. ``target_plan_id`` is required when ``role='critic'``.
|
| 155 |
+
"""
|
| 156 |
+
|
| 157 |
+
brain: Literal["epidemiology", "logistics", "governance"]
|
| 158 |
+
role: Literal["world_modeler", "planner", "critic"]
|
| 159 |
+
tick: int = Field(ge=0)
|
| 160 |
+
round: int = Field(ge=1, le=2, description="MVP cap: 1 or 2 only")
|
| 161 |
+
perception: PerceptionReport
|
| 162 |
+
prior_belief: Optional[BeliefState] = None
|
| 163 |
+
prior_plans: List[CandidatePlan] = Field(default_factory=list)
|
| 164 |
+
target_plan_id: Optional[str] = None
|
| 165 |
+
last_reward: float
|
| 166 |
+
recent_action_log_excerpt: List[ExecutedAction] = Field(default_factory=list)
|
| 167 |
+
|
| 168 |
+
|
| 169 |
+
class BrainLensedObservation(BaseModel):
|
| 170 |
+
"""Per-brain salience-mapped observation per Phase A §2 A1.
|
| 171 |
+
|
| 172 |
+
Lenses do not strip fields from the raw observation (Decision 13);
|
| 173 |
+
they project a salience map alongside it. ``derived_features`` lets
|
| 174 |
+
each brain pre-compute domain-specific scalars once and pass them
|
| 175 |
+
to all three of its LLM subagents without re-reading ``raw_obs``.
|
| 176 |
+
"""
|
| 177 |
+
|
| 178 |
+
brain: Literal["epidemiology", "logistics", "governance"]
|
| 179 |
+
raw_obs: CrisisworldcortexObservation
|
| 180 |
+
salient_field_ids: List[str] = Field(default_factory=list)
|
| 181 |
+
derived_features: Dict[str, float] = Field(default_factory=dict)
|
| 182 |
+
last_reward: float
|
| 183 |
+
|
| 184 |
+
|
| 185 |
# ============================================================================
|
| 186 |
# Brain output + Council decision
|
| 187 |
# ============================================================================
|
cortex/subagents/__init__.py
CHANGED
|
@@ -0,0 +1,27 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Cortex per-brain subagents (Session 9+).
|
| 2 |
+
|
| 3 |
+
Public surface:
|
| 4 |
+
- WorldModelerSubagent: emits BeliefState (LLM, router-callable).
|
| 5 |
+
- PlannerSubagent: emits CandidatePlan (LLM, router-callable).
|
| 6 |
+
- CriticSubagent: emits CriticReport (LLM, router-callable).
|
| 7 |
+
- perception_for: deterministic Python Perception function (Session 11+;
|
| 8 |
+
NOT router-callable per cortex/CLAUDE.md role-split binding).
|
| 9 |
+
- PROMPTS_DIR: directory holding the per-role SYS prompt templates.
|
| 10 |
+
|
| 11 |
+
Brain Executive (Python-only, NOT router-callable) lives in
|
| 12 |
+
``cortex/brains/_executive.py``.
|
| 13 |
+
"""
|
| 14 |
+
|
| 15 |
+
from ._base import PROMPTS_DIR
|
| 16 |
+
from .critic import CriticSubagent
|
| 17 |
+
from .perception import perception_for
|
| 18 |
+
from .planner import PlannerSubagent
|
| 19 |
+
from .world_modeler import WorldModelerSubagent
|
| 20 |
+
|
| 21 |
+
__all__ = [
|
| 22 |
+
"CriticSubagent",
|
| 23 |
+
"PROMPTS_DIR",
|
| 24 |
+
"PlannerSubagent",
|
| 25 |
+
"WorldModelerSubagent",
|
| 26 |
+
"perception_for",
|
| 27 |
+
]
|
cortex/subagents/_base.py
ADDED
|
@@ -0,0 +1,240 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Abstract base for the 3 LLM subagents (WorldModeler, Planner, Critic).
|
| 2 |
+
|
| 3 |
+
Phase A docs/CORTEX_ARCHITECTURE.md Decisions 1-8 + 62 lock the role split,
|
| 4 |
+
prompt-loading mechanism, retry-with-history semantics, empty fallback
|
| 5 |
+
shape, caller-id format, and TypeAdapter validation pattern. This base
|
| 6 |
+
class implements the shared mechanics; concrete subclasses pin the
|
| 7 |
+
role name, output type, prompt path, TypeAdapter, and USR builder.
|
| 8 |
+
"""
|
| 9 |
+
|
| 10 |
+
from __future__ import annotations
|
| 11 |
+
|
| 12 |
+
from abc import ABC, abstractmethod
|
| 13 |
+
from pathlib import Path
|
| 14 |
+
from typing import ClassVar, List, Optional, Protocol
|
| 15 |
+
|
| 16 |
+
from pydantic import BaseModel, TypeAdapter, ValidationError
|
| 17 |
+
|
| 18 |
+
from cortex.llm_client import ChatMessage, ChatResponse
|
| 19 |
+
from cortex.schemas import SubagentInput
|
| 20 |
+
from CrisisWorldCortex.models import ExecutedAction
|
| 21 |
+
|
| 22 |
+
# ============================================================================
|
| 23 |
+
# Module-level constants (loaded at import time)
|
| 24 |
+
# ============================================================================
|
| 25 |
+
|
| 26 |
+
PROMPTS_DIR: Path = Path(__file__).parent / "prompts"
|
| 27 |
+
"""Directory holding the per-role SYS prompt template files (Decision 4)."""
|
| 28 |
+
|
| 29 |
+
_RETRY_SNIPPET_MAX_CHARS: int = 200
|
| 30 |
+
"""Cap on the failed-response snippet included in the retry message
|
| 31 |
+
(Decision 8). 200 chars ~= 50 tokens; keeps retry overhead bounded."""
|
| 32 |
+
|
| 33 |
+
_BASE_RETRY_USER_TEMPLATE: str = (
|
| 34 |
+
"Your previous response failed to parse as JSON. The response was:\n"
|
| 35 |
+
"{snippet}\n\n"
|
| 36 |
+
"Emit ONLY valid JSON matching the schema specified in the system prompt. "
|
| 37 |
+
"No prose, no code fences."
|
| 38 |
+
)
|
| 39 |
+
|
| 40 |
+
_RECENT_ACTION_LOG_TAIL: int = 8
|
| 41 |
+
"""How many entries from ``recent_action_log_excerpt`` to render into the
|
| 42 |
+
USR summary. Matches the design-doc 8-deep history (M-FR-3)."""
|
| 43 |
+
|
| 44 |
+
|
| 45 |
+
# ============================================================================
|
| 46 |
+
# Duck-typed LLM client protocol (so tests can pass StubLLMClient)
|
| 47 |
+
# ============================================================================
|
| 48 |
+
|
| 49 |
+
|
| 50 |
+
class _LLMClientLike(Protocol):
|
| 51 |
+
"""Subset of ``cortex.llm_client.LLMClient`` that subagents call.
|
| 52 |
+
|
| 53 |
+
Production: ``LLMClient``. Tests: ``tests._helpers.llm_stub.StubLLMClient``.
|
| 54 |
+
"""
|
| 55 |
+
|
| 56 |
+
def chat(
|
| 57 |
+
self,
|
| 58 |
+
caller_id: str,
|
| 59 |
+
messages: List[ChatMessage],
|
| 60 |
+
max_tokens: Optional[int] = ...,
|
| 61 |
+
temperature: Optional[float] = ...,
|
| 62 |
+
) -> ChatResponse: ...
|
| 63 |
+
|
| 64 |
+
|
| 65 |
+
# ============================================================================
|
| 66 |
+
# Abstract base
|
| 67 |
+
# ============================================================================
|
| 68 |
+
|
| 69 |
+
|
| 70 |
+
class _LLMSubagent(ABC):
|
| 71 |
+
"""Shared run/retry/parse/empty-fallback skeleton for the 3 subagents.
|
| 72 |
+
|
| 73 |
+
Subclasses override the class-level vars below and implement
|
| 74 |
+
``_build_user_message`` + ``empty_fallback``.
|
| 75 |
+
"""
|
| 76 |
+
|
| 77 |
+
# --- Subclass class-level overrides -------------------------------------
|
| 78 |
+
_role_name: ClassVar[str] # one of: "world_modeler", "planner", "critic"
|
| 79 |
+
_output_type: ClassVar[type] # BeliefState / CandidatePlan / CriticReport
|
| 80 |
+
_system_prompt_filename: ClassVar[str] # e.g. "world_modeler.txt"
|
| 81 |
+
_SYSTEM_PROMPT_TEMPLATE: ClassVar[str] # populated by load_prompt() at module load
|
| 82 |
+
_ADAPTER: ClassVar[TypeAdapter] # populated at module load
|
| 83 |
+
|
| 84 |
+
# --- Construction --------------------------------------------------------
|
| 85 |
+
|
| 86 |
+
def __init__(self, llm_client: _LLMClientLike) -> None:
|
| 87 |
+
self._llm = llm_client
|
| 88 |
+
|
| 89 |
+
# --- Public surface ------------------------------------------------------
|
| 90 |
+
|
| 91 |
+
def run(self, input: SubagentInput, step_idx: int) -> BaseModel:
|
| 92 |
+
"""Call the LLM (with 1 retry), parse, return typed output or empty fallback.
|
| 93 |
+
|
| 94 |
+
Always returns a typed object - never ``None``. Decision 6: on
|
| 95 |
+
any failure (parse, retry-parse, LLM call exception) returns the
|
| 96 |
+
role-specific empty fallback.
|
| 97 |
+
"""
|
| 98 |
+
# Defensive: subclass enforces role-input alignment so harnesses
|
| 99 |
+
# don't accidentally route a Planner input through a Critic class.
|
| 100 |
+
assert input.role == self._role_name, (
|
| 101 |
+
f"SubagentInput.role={input.role!r} does not match "
|
| 102 |
+
f"{type(self).__name__}._role_name={self._role_name!r}"
|
| 103 |
+
)
|
| 104 |
+
|
| 105 |
+
sys_content = self._SYSTEM_PROMPT_TEMPLATE.format(
|
| 106 |
+
brain=input.brain,
|
| 107 |
+
target_plan_id=input.target_plan_id or "",
|
| 108 |
+
)
|
| 109 |
+
usr_content = self._build_user_message(input)
|
| 110 |
+
messages: List[ChatMessage] = [
|
| 111 |
+
ChatMessage(role="system", content=sys_content),
|
| 112 |
+
ChatMessage(role="user", content=usr_content),
|
| 113 |
+
]
|
| 114 |
+
caller_id = self._caller_id(input, step_idx)
|
| 115 |
+
|
| 116 |
+
# ---- Attempt 1 -----------------------------------------------------
|
| 117 |
+
first_response = self._safe_chat(caller_id, messages)
|
| 118 |
+
if first_response is None:
|
| 119 |
+
return self._empty_fallback_for(input)
|
| 120 |
+
parsed = self._try_parse(first_response.content)
|
| 121 |
+
if parsed is not None:
|
| 122 |
+
return parsed
|
| 123 |
+
|
| 124 |
+
# ---- Attempt 2 (retry with chat-history continuation) --------------
|
| 125 |
+
snippet = self._truncate_snippet(first_response.content)
|
| 126 |
+
retry_messages: List[ChatMessage] = [
|
| 127 |
+
*messages,
|
| 128 |
+
ChatMessage(role="assistant", content=first_response.content),
|
| 129 |
+
ChatMessage(
|
| 130 |
+
role="user",
|
| 131 |
+
content=_BASE_RETRY_USER_TEMPLATE.format(snippet=snippet),
|
| 132 |
+
),
|
| 133 |
+
]
|
| 134 |
+
retry_response = self._safe_chat(caller_id, retry_messages)
|
| 135 |
+
if retry_response is None:
|
| 136 |
+
return self._empty_fallback_for(input)
|
| 137 |
+
parsed_retry = self._try_parse(retry_response.content)
|
| 138 |
+
if parsed_retry is not None:
|
| 139 |
+
return parsed_retry
|
| 140 |
+
|
| 141 |
+
# ---- Both attempts failed - empty fallback -------------------------
|
| 142 |
+
return self._empty_fallback_for(input)
|
| 143 |
+
|
| 144 |
+
# --- Subclass extension points ------------------------------------------
|
| 145 |
+
|
| 146 |
+
@abstractmethod
|
| 147 |
+
def _build_user_message(self, input: SubagentInput) -> str:
|
| 148 |
+
"""Render the role-specific USR message body."""
|
| 149 |
+
|
| 150 |
+
@classmethod
|
| 151 |
+
@abstractmethod
|
| 152 |
+
def empty_fallback(cls, brain: str, target_plan_id: str = "") -> BaseModel:
|
| 153 |
+
"""Return the empty / no-signal output for this role.
|
| 154 |
+
|
| 155 |
+
Phase A Decision 6: confidence/severity = 0 and empty evidence/attacks
|
| 156 |
+
signal "no useful input from this subagent" to the Brain Executive.
|
| 157 |
+
"""
|
| 158 |
+
|
| 159 |
+
# --- Internal helpers ---------------------------------------------------
|
| 160 |
+
|
| 161 |
+
def _caller_id(self, input: SubagentInput, step_idx: int) -> str:
|
| 162 |
+
# Phase A Decision 7: cortex:<brain>:<role>:t<tick>:r<round>:s<step_idx>
|
| 163 |
+
return f"cortex:{input.brain}:{self._role_name}:t{input.tick}:r{input.round}:s{step_idx}"
|
| 164 |
+
|
| 165 |
+
def _safe_chat(self, caller_id: str, messages: List[ChatMessage]) -> Optional[ChatResponse]:
|
| 166 |
+
"""Call LLM; on exception, return None so caller can empty-fallback."""
|
| 167 |
+
try:
|
| 168 |
+
return self._llm.chat(caller_id=caller_id, messages=messages)
|
| 169 |
+
except Exception:
|
| 170 |
+
# Decision 6: LLM call failure folds into the same empty-fallback path
|
| 171 |
+
# as parse failure. Brain Executive sees a no-signal subagent.
|
| 172 |
+
return None
|
| 173 |
+
|
| 174 |
+
def _try_parse(self, content: str) -> Optional[BaseModel]:
|
| 175 |
+
"""Validate ``content`` as JSON via this role's TypeAdapter.
|
| 176 |
+
|
| 177 |
+
Strips common markdown code fences before validating since some
|
| 178 |
+
models wrap JSON in ```json ... ```.
|
| 179 |
+
"""
|
| 180 |
+
cleaned = _strip_code_fences(content.strip())
|
| 181 |
+
if not cleaned:
|
| 182 |
+
return None
|
| 183 |
+
try:
|
| 184 |
+
return self._ADAPTER.validate_json(cleaned)
|
| 185 |
+
except (ValidationError, ValueError):
|
| 186 |
+
return None
|
| 187 |
+
|
| 188 |
+
def _empty_fallback_for(self, input: SubagentInput) -> BaseModel:
|
| 189 |
+
return type(self).empty_fallback(
|
| 190 |
+
brain=input.brain,
|
| 191 |
+
target_plan_id=input.target_plan_id or "",
|
| 192 |
+
)
|
| 193 |
+
|
| 194 |
+
@staticmethod
|
| 195 |
+
def _truncate_snippet(content: str) -> str:
|
| 196 |
+
if len(content) <= _RETRY_SNIPPET_MAX_CHARS:
|
| 197 |
+
return content
|
| 198 |
+
return content[:_RETRY_SNIPPET_MAX_CHARS] + "..."
|
| 199 |
+
|
| 200 |
+
@staticmethod
|
| 201 |
+
def _format_action_log(log: List[ExecutedAction]) -> str:
|
| 202 |
+
"""M-FR-3 - render recent_action_log_excerpt as a compact text summary.
|
| 203 |
+
|
| 204 |
+
Format: ``"tick 4: deploy_resource accepted; tick 5: restrict_movement.strict rejected"``.
|
| 205 |
+
Capped at the most recent 8 entries.
|
| 206 |
+
"""
|
| 207 |
+
if not log:
|
| 208 |
+
return "(empty)"
|
| 209 |
+
items: List[str] = []
|
| 210 |
+
for ea in log[-_RECENT_ACTION_LOG_TAIL:]:
|
| 211 |
+
status = "accepted" if ea.accepted else "rejected"
|
| 212 |
+
kind = ea.action.kind
|
| 213 |
+
extra = ""
|
| 214 |
+
if kind == "restrict_movement":
|
| 215 |
+
extra = f".{getattr(ea.action, 'severity', '?')}"
|
| 216 |
+
items.append(f"tick {ea.tick}: {kind}{extra} {status}")
|
| 217 |
+
return "; ".join(items)
|
| 218 |
+
|
| 219 |
+
|
| 220 |
+
# ============================================================================
|
| 221 |
+
# Helpers (module-level)
|
| 222 |
+
# ============================================================================
|
| 223 |
+
|
| 224 |
+
|
| 225 |
+
def load_prompt(filename: str) -> str:
|
| 226 |
+
"""Load a SYS prompt template at module-load time (Decision 4)."""
|
| 227 |
+
return (PROMPTS_DIR / filename).read_text(encoding="utf-8")
|
| 228 |
+
|
| 229 |
+
|
| 230 |
+
def _strip_code_fences(s: str) -> str:
|
| 231 |
+
"""Remove leading ``` / ```json fence and trailing ``` if present."""
|
| 232 |
+
s = s.strip()
|
| 233 |
+
if not s.startswith("```"):
|
| 234 |
+
return s
|
| 235 |
+
lines = s.split("\n")
|
| 236 |
+
if lines[0].startswith("```"):
|
| 237 |
+
lines = lines[1:]
|
| 238 |
+
if lines and lines[-1].strip() == "```":
|
| 239 |
+
lines = lines[:-1]
|
| 240 |
+
return "\n".join(lines).strip()
|
cortex/subagents/critic.py
ADDED
|
@@ -0,0 +1,68 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Critic LLM subagent.
|
| 2 |
+
|
| 3 |
+
Phase A docs/CORTEX_ARCHITECTURE.md §9 Decision 3: SYS = critic role;
|
| 4 |
+
USR = perception + target plan + WM belief (M-FR-5). Critic emits prose
|
| 5 |
+
``CriticReport`` only; never proposes alternative actions.
|
| 6 |
+
"""
|
| 7 |
+
|
| 8 |
+
from __future__ import annotations
|
| 9 |
+
|
| 10 |
+
from typing import ClassVar, List
|
| 11 |
+
|
| 12 |
+
from pydantic import TypeAdapter
|
| 13 |
+
|
| 14 |
+
from cortex.schemas import CriticReport, SubagentInput
|
| 15 |
+
|
| 16 |
+
from ._base import _LLMSubagent, load_prompt
|
| 17 |
+
|
| 18 |
+
_CRITIC_ADAPTER: TypeAdapter[CriticReport] = TypeAdapter(CriticReport)
|
| 19 |
+
|
| 20 |
+
|
| 21 |
+
class CriticSubagent(_LLMSubagent):
|
| 22 |
+
"""LLM subagent that emits ``CriticReport`` for one brain per call."""
|
| 23 |
+
|
| 24 |
+
_role_name: ClassVar[str] = "critic"
|
| 25 |
+
_output_type: ClassVar[type] = CriticReport
|
| 26 |
+
_system_prompt_filename: ClassVar[str] = "critic.txt"
|
| 27 |
+
_SYSTEM_PROMPT_TEMPLATE: ClassVar[str] = load_prompt("critic.txt")
|
| 28 |
+
_ADAPTER: ClassVar[TypeAdapter] = _CRITIC_ADAPTER
|
| 29 |
+
|
| 30 |
+
def _build_user_message(self, input: SubagentInput) -> str:
|
| 31 |
+
sections: List[str] = []
|
| 32 |
+
sections.append(f"# Perception\n{input.perception.model_dump_json(indent=2)}")
|
| 33 |
+
# M-FR-5: target plan + WM belief, both as full JSON.
|
| 34 |
+
target_json = self._select_target_plan(input)
|
| 35 |
+
sections.append(f"# Target plan (id={input.target_plan_id})\n{target_json}")
|
| 36 |
+
if input.prior_belief is not None:
|
| 37 |
+
sections.append(f"# WM BeliefState\n{input.prior_belief.model_dump_json(indent=2)}")
|
| 38 |
+
return "\n\n".join(sections)
|
| 39 |
+
|
| 40 |
+
@staticmethod
|
| 41 |
+
def _select_target_plan(input: SubagentInput) -> str:
|
| 42 |
+
"""Render the target plan's JSON body for the USR.
|
| 43 |
+
|
| 44 |
+
Session 11's Brain Executive populates ``prior_plans`` from the
|
| 45 |
+
Planner's outputs and sets ``target_plan_id`` to identify which
|
| 46 |
+
plan the Critic should attack. Here we render the first plan
|
| 47 |
+
(or a placeholder if none) — Session 11 wires up id-based
|
| 48 |
+
lookup once plans carry ids.
|
| 49 |
+
"""
|
| 50 |
+
if not input.prior_plans:
|
| 51 |
+
return "(no target plan provided)"
|
| 52 |
+
return input.prior_plans[0].model_dump_json(indent=2)
|
| 53 |
+
|
| 54 |
+
@classmethod
|
| 55 |
+
def empty_fallback(cls, brain: str, target_plan_id: str = "") -> CriticReport:
|
| 56 |
+
# Phase A Decision 6: severity=0 + empty attacks signal "no
|
| 57 |
+
# critique". Brain Executive ignores this Critic's vote weight.
|
| 58 |
+
return CriticReport(
|
| 59 |
+
brain=brain,
|
| 60 |
+
target_plan_id=target_plan_id,
|
| 61 |
+
attacks=[],
|
| 62 |
+
missing_considerations=[],
|
| 63 |
+
would_change_mind_if=[],
|
| 64 |
+
severity=0.0,
|
| 65 |
+
)
|
| 66 |
+
|
| 67 |
+
def run(self, input: SubagentInput, step_idx: int) -> CriticReport: # type: ignore[override]
|
| 68 |
+
return super().run(input, step_idx) # type: ignore[return-value]
|
cortex/subagents/perception.py
ADDED
|
@@ -0,0 +1,192 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Perception subagent - deterministic Python; not router-callable.
|
| 2 |
+
|
| 3 |
+
Per cortex/CLAUDE.md binding: Perception is pure Python; no LLM calls.
|
| 4 |
+
Phase A Decisions 9 (V2 KeyError) + 63 (salient_signals cap at 5) +
|
| 5 |
+
M-FR-1 (pinned confidence per brain).
|
| 6 |
+
|
| 7 |
+
Perception runs ONCE per brain at tick start (not router-callable).
|
| 8 |
+
The Council Executive (Session 12) calls ``perception_for`` once per
|
| 9 |
+
brain at the start of each tick; the resulting ``PerceptionReport`` is
|
| 10 |
+
plumbed into all subsequent SubagentInputs for that brain in that tick.
|
| 11 |
+
"""
|
| 12 |
+
|
| 13 |
+
from __future__ import annotations
|
| 14 |
+
|
| 15 |
+
from typing import Callable, Dict, List
|
| 16 |
+
|
| 17 |
+
from cortex.schemas import EvidenceCitation, PerceptionReport
|
| 18 |
+
from CrisisWorldCortex.models import CrisisworldcortexObservation
|
| 19 |
+
|
| 20 |
+
_V2_BRAINS = frozenset({"communications", "equity"})
|
| 21 |
+
|
| 22 |
+
_SALIENT_SIGNALS_CAP = 5 # Phase A Decision 63 / OQ-2
|
| 23 |
+
|
| 24 |
+
_LOGISTICS_THRESHOLDS = {
|
| 25 |
+
"test_kits": 300,
|
| 26 |
+
"hospital_beds_free": 100,
|
| 27 |
+
"mobile_units": 5,
|
| 28 |
+
"vaccine_doses": 500,
|
| 29 |
+
}
|
| 30 |
+
|
| 31 |
+
_HOSPITAL_LOAD_ANOMALY_THRESHOLD = 0.6
|
| 32 |
+
|
| 33 |
+
|
| 34 |
+
def perception_for(brain: str, obs: CrisisworldcortexObservation) -> PerceptionReport:
|
| 35 |
+
"""Compute the per-brain Perception report.
|
| 36 |
+
|
| 37 |
+
Args:
|
| 38 |
+
brain: One of {"epidemiology", "logistics", "governance"}.
|
| 39 |
+
obs: The current tick's observation.
|
| 40 |
+
|
| 41 |
+
Raises:
|
| 42 |
+
KeyError: If ``brain`` is V2-deferred or unknown.
|
| 43 |
+
"""
|
| 44 |
+
helper = _PERCEPTION_REGISTRY.get(brain)
|
| 45 |
+
if helper is not None:
|
| 46 |
+
return helper(obs)
|
| 47 |
+
if brain in _V2_BRAINS:
|
| 48 |
+
raise KeyError(
|
| 49 |
+
f"V2 brain {brain!r} deferred per Phase A Decision 9; no MVP stub perception."
|
| 50 |
+
)
|
| 51 |
+
raise KeyError(f"unknown brain: {brain!r}")
|
| 52 |
+
|
| 53 |
+
|
| 54 |
+
def _epi_perception(obs: CrisisworldcortexObservation) -> PerceptionReport:
|
| 55 |
+
"""Epidemiology perception: top-cases regions + high-hospital-load anomalies."""
|
| 56 |
+
sorted_regions = sorted(obs.regions, key=lambda r: r.reported_cases_d_ago, reverse=True)
|
| 57 |
+
salient_signals: List[str] = []
|
| 58 |
+
evidence: List[EvidenceCitation] = []
|
| 59 |
+
|
| 60 |
+
for r in sorted_regions[:3]:
|
| 61 |
+
if r.reported_cases_d_ago > 0:
|
| 62 |
+
salient_signals.append(f"{r.region}: cases={r.reported_cases_d_ago}")
|
| 63 |
+
evidence.append(
|
| 64 |
+
EvidenceCitation(
|
| 65 |
+
source="telemetry",
|
| 66 |
+
ref=f"{r.region}.reported_cases_d_ago",
|
| 67 |
+
excerpt=str(r.reported_cases_d_ago),
|
| 68 |
+
)
|
| 69 |
+
)
|
| 70 |
+
|
| 71 |
+
if not salient_signals and obs.regions:
|
| 72 |
+
# Fallback: cite the first region so we have at least one signal
|
| 73 |
+
r = obs.regions[0]
|
| 74 |
+
salient_signals.append(f"{r.region}: cases={r.reported_cases_d_ago}")
|
| 75 |
+
evidence.append(
|
| 76 |
+
EvidenceCitation(
|
| 77 |
+
source="telemetry",
|
| 78 |
+
ref=f"{r.region}.reported_cases_d_ago",
|
| 79 |
+
excerpt=str(r.reported_cases_d_ago),
|
| 80 |
+
)
|
| 81 |
+
)
|
| 82 |
+
|
| 83 |
+
salient_signals = salient_signals[:_SALIENT_SIGNALS_CAP]
|
| 84 |
+
|
| 85 |
+
anomalies = [
|
| 86 |
+
f"{r.region}: hospital_load={r.hospital_load:.2f}"
|
| 87 |
+
for r in obs.regions
|
| 88 |
+
if r.hospital_load > _HOSPITAL_LOAD_ANOMALY_THRESHOLD
|
| 89 |
+
]
|
| 90 |
+
|
| 91 |
+
return PerceptionReport(
|
| 92 |
+
brain="epidemiology",
|
| 93 |
+
salient_signals=salient_signals,
|
| 94 |
+
anomalies=anomalies,
|
| 95 |
+
# M-FR-1: telemetry is delayed and noisy per mm.md; pinned proxy
|
| 96 |
+
confidence=0.7,
|
| 97 |
+
evidence=evidence,
|
| 98 |
+
)
|
| 99 |
+
|
| 100 |
+
|
| 101 |
+
def _logistics_perception(obs: CrisisworldcortexObservation) -> PerceptionReport:
|
| 102 |
+
"""Logistics perception: low-resource flags + depleted-resource anomalies."""
|
| 103 |
+
res = obs.resources
|
| 104 |
+
salient_signals: List[str] = []
|
| 105 |
+
evidence: List[EvidenceCitation] = []
|
| 106 |
+
|
| 107 |
+
for resource_name, threshold in _LOGISTICS_THRESHOLDS.items():
|
| 108 |
+
value = getattr(res, resource_name)
|
| 109 |
+
if value < threshold:
|
| 110 |
+
salient_signals.append(f"{resource_name} low: {value}")
|
| 111 |
+
evidence.append(
|
| 112 |
+
EvidenceCitation(
|
| 113 |
+
source="resource",
|
| 114 |
+
ref=f"resources.{resource_name}",
|
| 115 |
+
excerpt=str(value),
|
| 116 |
+
)
|
| 117 |
+
)
|
| 118 |
+
|
| 119 |
+
salient_signals = salient_signals[:_SALIENT_SIGNALS_CAP]
|
| 120 |
+
|
| 121 |
+
anomalies = []
|
| 122 |
+
for resource_name in _LOGISTICS_THRESHOLDS:
|
| 123 |
+
if getattr(res, resource_name) == 0:
|
| 124 |
+
anomalies.append(f"{resource_name}: depleted")
|
| 125 |
+
|
| 126 |
+
return PerceptionReport(
|
| 127 |
+
brain="logistics",
|
| 128 |
+
salient_signals=salient_signals,
|
| 129 |
+
anomalies=anomalies,
|
| 130 |
+
# M-FR-1: resource counts are deterministic, no telemetry noise
|
| 131 |
+
confidence=1.0,
|
| 132 |
+
evidence=evidence,
|
| 133 |
+
)
|
| 134 |
+
|
| 135 |
+
|
| 136 |
+
def _governance_perception(obs: CrisisworldcortexObservation) -> PerceptionReport:
|
| 137 |
+
"""Governance perception: active restrictions + legal constraints + about-to-expire anomalies."""
|
| 138 |
+
salient_signals: List[str] = []
|
| 139 |
+
evidence: List[EvidenceCitation] = []
|
| 140 |
+
|
| 141 |
+
for restr in obs.active_restrictions:
|
| 142 |
+
salient_signals.append(f"{restr.region}: {restr.severity} ({restr.ticks_remaining}t)")
|
| 143 |
+
evidence.append(
|
| 144 |
+
EvidenceCitation(
|
| 145 |
+
source="policy",
|
| 146 |
+
ref=f"active_restrictions.{restr.region}",
|
| 147 |
+
excerpt=f"{restr.severity}@{restr.ticks_remaining}",
|
| 148 |
+
)
|
| 149 |
+
)
|
| 150 |
+
|
| 151 |
+
for lc in obs.legal_constraints:
|
| 152 |
+
salient_signals.append(f"legal: {lc.rule_id} blocks {lc.blocked_action}")
|
| 153 |
+
evidence.append(
|
| 154 |
+
EvidenceCitation(
|
| 155 |
+
source="policy",
|
| 156 |
+
ref=f"legal_constraints.{lc.rule_id}",
|
| 157 |
+
excerpt=lc.blocked_action,
|
| 158 |
+
)
|
| 159 |
+
)
|
| 160 |
+
|
| 161 |
+
salient_signals = salient_signals[:_SALIENT_SIGNALS_CAP]
|
| 162 |
+
|
| 163 |
+
has_recent_escalate_national = any(
|
| 164 |
+
ea.accepted
|
| 165 |
+
and ea.action.kind == "escalate"
|
| 166 |
+
and getattr(ea.action, "to_authority", None) == "national"
|
| 167 |
+
for ea in obs.recent_action_log
|
| 168 |
+
)
|
| 169 |
+
anomalies = []
|
| 170 |
+
for restr in obs.active_restrictions:
|
| 171 |
+
if (
|
| 172 |
+
restr.severity == "strict"
|
| 173 |
+
and restr.ticks_remaining <= 1
|
| 174 |
+
and not has_recent_escalate_national
|
| 175 |
+
):
|
| 176 |
+
anomalies.append(f"{restr.region}: strict expiring without escalation")
|
| 177 |
+
|
| 178 |
+
return PerceptionReport(
|
| 179 |
+
brain="governance",
|
| 180 |
+
salient_signals=salient_signals,
|
| 181 |
+
anomalies=anomalies,
|
| 182 |
+
# M-FR-1: policy state is deterministic
|
| 183 |
+
confidence=1.0,
|
| 184 |
+
evidence=evidence,
|
| 185 |
+
)
|
| 186 |
+
|
| 187 |
+
|
| 188 |
+
_PERCEPTION_REGISTRY: Dict[str, Callable[[CrisisworldcortexObservation], PerceptionReport]] = {
|
| 189 |
+
"epidemiology": _epi_perception,
|
| 190 |
+
"logistics": _logistics_perception,
|
| 191 |
+
"governance": _governance_perception,
|
| 192 |
+
}
|
cortex/subagents/planner.py
ADDED
|
@@ -0,0 +1,61 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Planner LLM subagent.
|
| 2 |
+
|
| 3 |
+
Phase A docs/CORTEX_ARCHITECTURE.md §9 Decision 2: SYS = role + action
|
| 4 |
+
schema (B1's shape); USR = perception + WM BeliefState (full JSON if
|
| 5 |
+
provided per M-FR-4) + last_reward.
|
| 6 |
+
"""
|
| 7 |
+
|
| 8 |
+
from __future__ import annotations
|
| 9 |
+
|
| 10 |
+
from typing import ClassVar, List
|
| 11 |
+
|
| 12 |
+
from pydantic import TypeAdapter
|
| 13 |
+
|
| 14 |
+
from cortex.schemas import CandidatePlan, SubagentInput
|
| 15 |
+
from CrisisWorldCortex.models import NoOp
|
| 16 |
+
|
| 17 |
+
from ._base import _LLMSubagent, load_prompt
|
| 18 |
+
|
| 19 |
+
_PLAN_ADAPTER: TypeAdapter[CandidatePlan] = TypeAdapter(CandidatePlan)
|
| 20 |
+
|
| 21 |
+
|
| 22 |
+
class PlannerSubagent(_LLMSubagent):
|
| 23 |
+
"""LLM subagent that emits ``CandidatePlan`` for one brain per call."""
|
| 24 |
+
|
| 25 |
+
_role_name: ClassVar[str] = "planner"
|
| 26 |
+
_output_type: ClassVar[type] = CandidatePlan
|
| 27 |
+
_system_prompt_filename: ClassVar[str] = "planner.txt"
|
| 28 |
+
_SYSTEM_PROMPT_TEMPLATE: ClassVar[str] = load_prompt("planner.txt")
|
| 29 |
+
_ADAPTER: ClassVar[TypeAdapter] = _PLAN_ADAPTER
|
| 30 |
+
|
| 31 |
+
def _build_user_message(self, input: SubagentInput) -> str:
|
| 32 |
+
sections: List[str] = []
|
| 33 |
+
sections.append(f"# Perception\n{input.perception.model_dump_json(indent=2)}")
|
| 34 |
+
if input.prior_belief is not None:
|
| 35 |
+
sections.append(
|
| 36 |
+
"# BeliefState (from this brain's WorldModeler)\n"
|
| 37 |
+
f"{input.prior_belief.model_dump_json(indent=2)}"
|
| 38 |
+
)
|
| 39 |
+
sections.append(f"# Last tick reward: {input.last_reward}")
|
| 40 |
+
sections.append(
|
| 41 |
+
f"# Recent action log: {self._format_action_log(input.recent_action_log_excerpt)}"
|
| 42 |
+
)
|
| 43 |
+
return "\n\n".join(sections)
|
| 44 |
+
|
| 45 |
+
@classmethod
|
| 46 |
+
def empty_fallback(cls, brain: str, target_plan_id: str = "") -> CandidatePlan:
|
| 47 |
+
# Phase A Decision 6: NoOp + confidence=0 means "no signal". The
|
| 48 |
+
# Brain Executive's argmax(expected_value * confidence) picks any
|
| 49 |
+
# non-empty plan over this one.
|
| 50 |
+
return CandidatePlan(
|
| 51 |
+
action_sketch="(empty: planner failed to produce a parseable plan)",
|
| 52 |
+
expected_outer_action=NoOp(),
|
| 53 |
+
expected_value=0.0,
|
| 54 |
+
cost=0.0,
|
| 55 |
+
assumptions=[],
|
| 56 |
+
falsifiers=[],
|
| 57 |
+
confidence=0.0,
|
| 58 |
+
)
|
| 59 |
+
|
| 60 |
+
def run(self, input: SubagentInput, step_idx: int) -> CandidatePlan: # type: ignore[override]
|
| 61 |
+
return super().run(input, step_idx) # type: ignore[return-value]
|
cortex/subagents/prompts/critic.txt
ADDED
|
@@ -0,0 +1,19 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
You are the {brain} brain's Critic in the CrisisWorldCortex multi-brain agent.
|
| 2 |
+
|
| 3 |
+
Your job: read a CandidatePlan from this brain's Planner and identify what is wrong with it. Emit a CriticReport as JSON.
|
| 4 |
+
|
| 5 |
+
You write prose critique only. Do NOT propose alternative actions; that is the Planner's role.
|
| 6 |
+
|
| 7 |
+
Output strict JSON matching this schema:
|
| 8 |
+
{{
|
| 9 |
+
"brain": "{brain}",
|
| 10 |
+
"target_plan_id": "{target_plan_id}",
|
| 11 |
+
"attacks": ["<one sentence per concrete attack>"],
|
| 12 |
+
"missing_considerations": ["<one sentence per missed factor>"],
|
| 13 |
+
"would_change_mind_if": ["<observation or fact that would soften your critique>"],
|
| 14 |
+
"severity": <float in [0,1]; 0=plan is fine, 1=do not execute this plan>
|
| 15 |
+
}}
|
| 16 |
+
|
| 17 |
+
Use prose strings inside the lists. Do NOT emit JSON action variants.
|
| 18 |
+
|
| 19 |
+
Emit ONLY the JSON. No prose outside the JSON, no code fences.
|
cortex/subagents/prompts/planner.txt
ADDED
|
@@ -0,0 +1,26 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
You are the {brain} brain's Planner in the CrisisWorldCortex multi-brain agent.
|
| 2 |
+
|
| 3 |
+
Your job: read the perception summary and prior BeliefState, then emit ONE candidate plan as JSON.
|
| 4 |
+
|
| 5 |
+
Output strict JSON matching this schema:
|
| 6 |
+
{{
|
| 7 |
+
"action_sketch": "<one sentence describing the plan>",
|
| 8 |
+
"expected_outer_action": <one of the action variants below>,
|
| 9 |
+
"expected_value": <float; estimated reward delta for this tick>,
|
| 10 |
+
"cost": <float; relative cost of this plan>,
|
| 11 |
+
"assumptions": ["<assumption 1>"],
|
| 12 |
+
"falsifiers": ["<observation that would refute this plan>"],
|
| 13 |
+
"confidence": <float in [0,1]>
|
| 14 |
+
}}
|
| 15 |
+
|
| 16 |
+
Action variants for expected_outer_action (exactly one kind):
|
| 17 |
+
- {{"kind":"deploy_resource", "region":"R1|R2|R3|R4", "resource_type":"test_kits|hospital_beds|mobile_units|vaccine_doses", "quantity": <int >= 0>}}
|
| 18 |
+
- {{"kind":"request_data", "region":"<id>", "data_type":"case_survey|hospital_audit|compliance_check"}}
|
| 19 |
+
- {{"kind":"restrict_movement", "region":"<id>", "severity":"none|light|moderate|strict"}}
|
| 20 |
+
- {{"kind":"escalate", "to_authority":"regional|national"}}
|
| 21 |
+
- {{"kind":"reallocate_budget", "from_resource":"<resource_type>", "to_resource":"<resource_type>", "amount": <int >= 0>}}
|
| 22 |
+
- {{"kind":"no_op"}}
|
| 23 |
+
|
| 24 |
+
Strict severity may require a prior escalate(national) - check legal_constraints.
|
| 25 |
+
|
| 26 |
+
Emit ONLY the JSON. No prose, no code fences.
|
cortex/subagents/prompts/world_modeler.txt
ADDED
|
@@ -0,0 +1,32 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
You are the {brain} brain's World Modeler in the CrisisWorldCortex multi-brain agent.
|
| 2 |
+
|
| 3 |
+
Your job: read the perception summary plus prior belief (round 2 only), and emit a BeliefState as JSON.
|
| 4 |
+
|
| 5 |
+
A BeliefState describes what you think the latent epidemiological / logistical / governance state is for each region, citing evidence from observed telemetry, resources, policy state, or recent actions.
|
| 6 |
+
|
| 7 |
+
Output strict JSON matching this schema:
|
| 8 |
+
{{
|
| 9 |
+
"brain": "{brain}",
|
| 10 |
+
"latent_estimates": {{
|
| 11 |
+
"<region_id>": {{
|
| 12 |
+
"estimated_infection_rate": <float in [0,1]>,
|
| 13 |
+
"estimated_r_effective": <float >= 0>,
|
| 14 |
+
"estimated_compliance": <float in [0,1]>,
|
| 15 |
+
"confidence_intervals": {{}}
|
| 16 |
+
}}
|
| 17 |
+
}},
|
| 18 |
+
"hypotheses": [
|
| 19 |
+
{{"label": "<short>", "weight": <float in [0,1]>, "explanation": "<one sentence>"}}
|
| 20 |
+
],
|
| 21 |
+
"uncertainty": <float in [0,1]>,
|
| 22 |
+
"reducible_by_more_thought": <float in [0,1]>,
|
| 23 |
+
"evidence": [
|
| 24 |
+
{{"source": "telemetry|resource|policy|action_log|belief|memory",
|
| 25 |
+
"ref": "<concise pointer like region=R2.hospital_load@tick=7>",
|
| 26 |
+
"excerpt": "<the value or text>"}}
|
| 27 |
+
]
|
| 28 |
+
}}
|
| 29 |
+
|
| 30 |
+
Cite at least 2 EvidenceCitations. Uncited claims zero your protocol-integrity reward.
|
| 31 |
+
|
| 32 |
+
Emit ONLY the JSON. No prose, no code fences.
|
cortex/subagents/world_modeler.py
ADDED
|
@@ -0,0 +1,61 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""WorldModeler LLM subagent.
|
| 2 |
+
|
| 3 |
+
Phase A docs/CORTEX_ARCHITECTURE.md §9 Decision 1: SYS = role + schema;
|
| 4 |
+
USR = perception + last_reward + recent_action_log_excerpt (with prior
|
| 5 |
+
BeliefState in round 2 per Decision 62).
|
| 6 |
+
"""
|
| 7 |
+
|
| 8 |
+
from __future__ import annotations
|
| 9 |
+
|
| 10 |
+
from typing import ClassVar, List
|
| 11 |
+
|
| 12 |
+
from pydantic import TypeAdapter
|
| 13 |
+
|
| 14 |
+
from cortex.schemas import BeliefState, SubagentInput
|
| 15 |
+
|
| 16 |
+
from ._base import _LLMSubagent, load_prompt
|
| 17 |
+
|
| 18 |
+
_BELIEF_ADAPTER: TypeAdapter[BeliefState] = TypeAdapter(BeliefState)
|
| 19 |
+
"""Module-level constant per Phase A Decision 5 (encapsulation; avoid
|
| 20 |
+
circular imports through the package init)."""
|
| 21 |
+
|
| 22 |
+
|
| 23 |
+
class WorldModelerSubagent(_LLMSubagent):
|
| 24 |
+
"""LLM subagent that emits ``BeliefState`` for one brain per call."""
|
| 25 |
+
|
| 26 |
+
_role_name: ClassVar[str] = "world_modeler"
|
| 27 |
+
_output_type: ClassVar[type] = BeliefState
|
| 28 |
+
_system_prompt_filename: ClassVar[str] = "world_modeler.txt"
|
| 29 |
+
_SYSTEM_PROMPT_TEMPLATE: ClassVar[str] = load_prompt("world_modeler.txt")
|
| 30 |
+
_ADAPTER: ClassVar[TypeAdapter] = _BELIEF_ADAPTER
|
| 31 |
+
|
| 32 |
+
def _build_user_message(self, input: SubagentInput) -> str:
|
| 33 |
+
sections: List[str] = []
|
| 34 |
+
sections.append(f"# Perception\n{input.perception.model_dump_json(indent=2)}")
|
| 35 |
+
if input.prior_belief is not None:
|
| 36 |
+
sections.append(
|
| 37 |
+
"# Prior BeliefState (round 1 result)\n"
|
| 38 |
+
f"{input.prior_belief.model_dump_json(indent=2)}"
|
| 39 |
+
)
|
| 40 |
+
sections.append(f"# Last tick reward: {input.last_reward}")
|
| 41 |
+
sections.append(
|
| 42 |
+
f"# Recent action log: {self._format_action_log(input.recent_action_log_excerpt)}"
|
| 43 |
+
)
|
| 44 |
+
return "\n\n".join(sections)
|
| 45 |
+
|
| 46 |
+
@classmethod
|
| 47 |
+
def empty_fallback(cls, brain: str, target_plan_id: str = "") -> BeliefState:
|
| 48 |
+
# Phase A Decision 6 + Decision 62: empty BeliefState as the
|
| 49 |
+
# honest "no signal" state. uncertainty=1.0, no evidence -> r_proto = 0.
|
| 50 |
+
return BeliefState(
|
| 51 |
+
brain=brain,
|
| 52 |
+
latent_estimates={},
|
| 53 |
+
hypotheses=[],
|
| 54 |
+
uncertainty=1.0,
|
| 55 |
+
reducible_by_more_thought=0.0,
|
| 56 |
+
evidence=[],
|
| 57 |
+
)
|
| 58 |
+
|
| 59 |
+
# Narrow run() return type for callers (refinement #1).
|
| 60 |
+
def run(self, input: SubagentInput, step_idx: int) -> BeliefState: # type: ignore[override]
|
| 61 |
+
return super().run(input, step_idx) # type: ignore[return-value]
|
demo/CLAUDE.md
CHANGED
|
@@ -1,45 +1,45 @@
|
|
| 1 |
-
# demo/CLAUDE.md
|
| 2 |
-
|
| 3 |
-
Replay-only visualization. Live demos fail under judging pressure; ship canned scenarios.
|
| 4 |
-
|
| 5 |
-
## Belongs here
|
| 6 |
-
|
| 7 |
-
- `visualizer/trace_renderer.py` — renders a JSON trace as a "council in action" view.
|
| 8 |
-
- `visualizer/reward_curve_plot.py` — plots reward curves from training logs.
|
| 9 |
-
- `demo_scenarios/*.json` — pre-recorded trajectories for the pitch (e.g. `scenario_flat_fails.json`, `scenario_cortex_holds_dissent.json`).
|
| 10 |
-
|
| 11 |
-
## Does not belong here
|
| 12 |
-
|
| 13 |
-
Live agent execution (record offline, replay here). Training logic. Graders.
|
| 14 |
-
|
| 15 |
-
## Allowed imports
|
| 16 |
-
|
| 17 |
-
- `cortex.schemas` — typed parse of trace JSON. Types only, no logic.
|
| 18 |
-
- stdlib + plotting libs (matplotlib / plotly).
|
| 19 |
-
|
| 20 |
-
## Forbidden imports
|
| 21 |
-
|
| 22 |
-
- `server/*`, `training/*`, `baselines/*`.
|
| 23 |
-
- `cortex.council`, `cortex.routing_policy` — if you need to re-run the agent, do it offline and ship a new JSON.
|
| 24 |
-
|
| 25 |
-
## Binding contracts
|
| 26 |
-
|
| 27 |
-
- Every JSON scenario conforms to `cortex.schemas.Trajectory`.
|
| 28 |
-
- Rendering is deterministic: same JSON → same output, modulo timestamps.
|
| 29 |
-
- The pitch-demo scenario must showcase B2 overcommit/misallocate vs Cortex dissent-preservation (design §27).
|
| 30 |
-
- A pre-recorded demo video (MP4) lives alongside the JSON scenarios as the live-demo fallback.
|
| 31 |
-
|
| 32 |
-
## Public APIs (owned here)
|
| 33 |
-
|
| 34 |
-
- `render_trace(json_path: str, out_path: str) -> None`
|
| 35 |
-
- `plot_reward_curves(log_paths: list[str], out_path: str) -> None`
|
| 36 |
-
|
| 37 |
-
## Testing requirements
|
| 38 |
-
|
| 39 |
-
- Each committed JSON scenario parses into a `Trajectory` without error.
|
| 40 |
-
- `render_trace` produces a non-empty output file for each scenario.
|
| 41 |
-
|
| 42 |
-
## Common failure modes
|
| 43 |
-
|
| 44 |
-
- Live re-run during the demo — network/Colab flakiness kills the pitch. Replay only.
|
| 45 |
-
- Renderer depending on a `cortex.council` instance — import breaks when Cortex API shifts. Keep read-only on types.
|
|
|
|
| 1 |
+
# demo/CLAUDE.md
|
| 2 |
+
|
| 3 |
+
Replay-only visualization. Live demos fail under judging pressure; ship canned scenarios.
|
| 4 |
+
|
| 5 |
+
## Belongs here
|
| 6 |
+
|
| 7 |
+
- `visualizer/trace_renderer.py` — renders a JSON trace as a "council in action" view.
|
| 8 |
+
- `visualizer/reward_curve_plot.py` — plots reward curves from training logs.
|
| 9 |
+
- `demo_scenarios/*.json` — pre-recorded trajectories for the pitch (e.g. `scenario_flat_fails.json`, `scenario_cortex_holds_dissent.json`).
|
| 10 |
+
|
| 11 |
+
## Does not belong here
|
| 12 |
+
|
| 13 |
+
Live agent execution (record offline, replay here). Training logic. Graders.
|
| 14 |
+
|
| 15 |
+
## Allowed imports
|
| 16 |
+
|
| 17 |
+
- `cortex.schemas` — typed parse of trace JSON. Types only, no logic.
|
| 18 |
+
- stdlib + plotting libs (matplotlib / plotly).
|
| 19 |
+
|
| 20 |
+
## Forbidden imports
|
| 21 |
+
|
| 22 |
+
- `server/*`, `training/*`, `baselines/*`.
|
| 23 |
+
- `cortex.council`, `cortex.routing_policy` — if you need to re-run the agent, do it offline and ship a new JSON.
|
| 24 |
+
|
| 25 |
+
## Binding contracts
|
| 26 |
+
|
| 27 |
+
- Every JSON scenario conforms to `cortex.schemas.Trajectory`.
|
| 28 |
+
- Rendering is deterministic: same JSON → same output, modulo timestamps.
|
| 29 |
+
- The pitch-demo scenario must showcase B2 overcommit/misallocate vs Cortex dissent-preservation (design §27).
|
| 30 |
+
- A pre-recorded demo video (MP4) lives alongside the JSON scenarios as the live-demo fallback.
|
| 31 |
+
|
| 32 |
+
## Public APIs (owned here)
|
| 33 |
+
|
| 34 |
+
- `render_trace(json_path: str, out_path: str) -> None`
|
| 35 |
+
- `plot_reward_curves(log_paths: list[str], out_path: str) -> None`
|
| 36 |
+
|
| 37 |
+
## Testing requirements
|
| 38 |
+
|
| 39 |
+
- Each committed JSON scenario parses into a `Trajectory` without error.
|
| 40 |
+
- `render_trace` produces a non-empty output file for each scenario.
|
| 41 |
+
|
| 42 |
+
## Common failure modes
|
| 43 |
+
|
| 44 |
+
- Live re-run during the demo — network/Colab flakiness kills the pitch. Replay only.
|
| 45 |
+
- Renderer depending on a `cortex.council` instance — import breaks when Cortex API shifts. Keep read-only on types.
|
docs/CORTEX_ARCHITECTURE.md
CHANGED
|
@@ -705,7 +705,7 @@ Decisions are grouped by layer. Each entry: **decision** / **rationale** /
|
|
| 705 |
|
| 706 |
19. **Reasoning summary: a 1–2 sentence string from the Planner's `action_sketch`.** / Fits the 400-char `BrainRecommendation.reasoning_summary` cap. / Considered an LLM call to produce a summary; rejected — Brain Executive must be Python-only per cortex/CLAUDE.md.
|
| 707 |
|
| 708 |
-
20. **`evidence` field on `BrainRecommendation` = union of all `EvidenceCitation` lists from `BeliefState`, `CandidatePlan`, `CriticReport`.** / Ensures the council sees the brain's full evidence chain. / Considered Critic only; rejected — claims with no upstream evidence get zeroed `r_proto`.
|
| 709 |
|
| 710 |
21. **Brain identifier strings: `"epidemiology"`, `"logistics"`, `"governance"` (lowercase, full word).** / Readable and grep-friendly. / Considered abbreviations (epi, log, gov); rejected — log-grep collisions.
|
| 711 |
|
|
|
|
| 705 |
|
| 706 |
19. **Reasoning summary: a 1–2 sentence string from the Planner's `action_sketch`.** / Fits the 400-char `BrainRecommendation.reasoning_summary` cap. / Considered an LLM call to produce a summary; rejected — Brain Executive must be Python-only per cortex/CLAUDE.md.
|
| 707 |
|
| 708 |
+
20. **`evidence` field on `BrainRecommendation` = union of all `EvidenceCitation` lists from `BeliefState`, `CandidatePlan`, `CriticReport`.** / Ensures the council sees the brain's full evidence chain. / Considered Critic only; rejected — claims with no upstream evidence get zeroed `r_proto`. **(Session 11 implementation note — M-FR-3)** Implementation reads evidence from `PerceptionReport.evidence` + `BeliefState.evidence` only, since `CandidatePlan` and `CriticReport` schemas (Session 9) carry no `evidence` field. Adding evidence fields to those schemas was rejected as schema-churn risk; the perception+beliefs union captures the actionable evidence chain since plans/critics derive from beliefs. See `cortex/brains/_executive.py:aggregate_brain_outputs`.
|
| 709 |
|
| 710 |
21. **Brain identifier strings: `"epidemiology"`, `"logistics"`, `"governance"` (lowercase, full word).** / Readable and grep-friendly. / Considered abbreviations (epi, log, gov); rejected — log-grep collisions.
|
| 711 |
|
inference.py
CHANGED
|
@@ -40,10 +40,15 @@ from __future__ import annotations
|
|
| 40 |
import os
|
| 41 |
import sys
|
| 42 |
from dataclasses import dataclass
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 43 |
from typing import Any, Dict, List, Optional
|
| 44 |
|
| 45 |
-
from baselines.flat_agent import B1FlatAgent, B1StepEvent
|
| 46 |
-
from cortex.llm_client import LLMClient
|
| 47 |
from CrisisWorldCortex.models import OuterActionPayload
|
| 48 |
from CrisisWorldCortex.server.graders import terminal_bonus
|
| 49 |
from CrisisWorldCortex.server.simulator import WorldState
|
|
@@ -62,8 +67,8 @@ DEFAULT_MODEL = "Qwen/Qwen2.5-72B-Instruct"
|
|
| 62 |
# distinct seeds per task for cross-episode reproducibility.
|
| 63 |
TASK_CONFIGS: List[dict] = [
|
| 64 |
{"task_name": "outbreak_easy", "seed": 0, "max_ticks": 12},
|
| 65 |
-
{"task_name": "outbreak_medium", "seed": 1, "max_ticks": 12},
|
| 66 |
-
{"task_name": "outbreak_hard", "seed": 2, "max_ticks": 12},
|
| 67 |
]
|
| 68 |
|
| 69 |
# Score-clamp bounds keep .3f formatting strictly inside (0, 1) so the
|
|
|
|
| 40 |
import os
|
| 41 |
import sys
|
| 42 |
from dataclasses import dataclass
|
| 43 |
+
try:
|
| 44 |
+
from dotenv import load_dotenv
|
| 45 |
+
load_dotenv()
|
| 46 |
+
except ImportError:
|
| 47 |
+
pass
|
| 48 |
from typing import Any, Dict, List, Optional
|
| 49 |
|
| 50 |
+
from CrisisWorldCortex.baselines.flat_agent import B1FlatAgent, B1StepEvent
|
| 51 |
+
from CrisisWorldCortex.cortex.llm_client import LLMClient
|
| 52 |
from CrisisWorldCortex.models import OuterActionPayload
|
| 53 |
from CrisisWorldCortex.server.graders import terminal_bonus
|
| 54 |
from CrisisWorldCortex.server.simulator import WorldState
|
|
|
|
| 67 |
# distinct seeds per task for cross-episode reproducibility.
|
| 68 |
TASK_CONFIGS: List[dict] = [
|
| 69 |
{"task_name": "outbreak_easy", "seed": 0, "max_ticks": 12},
|
| 70 |
+
# {"task_name": "outbreak_medium", "seed": 1, "max_ticks": 12},
|
| 71 |
+
# {"task_name": "outbreak_hard", "seed": 2, "max_ticks": 12},
|
| 72 |
]
|
| 73 |
|
| 74 |
# Score-clamp bounds keep .3f formatting strictly inside (0, 1) so the
|
notebooks/train_b1_grpo.ipynb
ADDED
|
@@ -0,0 +1,496 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"cells": [
|
| 3 |
+
{
|
| 4 |
+
"cell_type": "markdown",
|
| 5 |
+
"metadata": {},
|
| 6 |
+
"source": [
|
| 7 |
+
"# B1 GRPO training on Qwen3-1.7B (Workstream B Phase 3)\n",
|
| 8 |
+
"\n",
|
| 9 |
+
"Trains the **B1 flat-agent baseline** with [Unsloth](https://docs.unsloth.ai/get-started/reinforcement-learning-rl-guide) + [TRL GRPO](https://huggingface.co/docs/trl/main/en/grpo_trainer) against the deployed CrisisWorldCortex HF Space env.\n",
|
| 10 |
+
"\n",
|
| 11 |
+
"**One-shot run:** `Runtime → Run all` on a fresh Colab T4. 300 training steps, ~30 minutes wall-clock. Saves the trained LoRA adapter to your HF Hub at the end.\n",
|
| 12 |
+
"\n",
|
| 13 |
+
"**Reward source:** the Phase-1-fixed `outer_reward` (range `[-1.0, 1.0]`, signal-quality gates passed). Single-step GRPO — each rollout = one env reset + one env step.\n",
|
| 14 |
+
"\n",
|
| 15 |
+
"**Prereqs:**\n",
|
| 16 |
+
"1. Colab Secrets has `HF_TOKEN` set (Tools → Secrets, name = `HF_TOKEN`, value = your `hf_xxx` token with write access).\n",
|
| 17 |
+
"2. The HF Space `Angshuman28/CrisisWorldCortex` is running with the post-Phase-1 reward.\n"
|
| 18 |
+
]
|
| 19 |
+
},
|
| 20 |
+
{
|
| 21 |
+
"cell_type": "markdown",
|
| 22 |
+
"metadata": {},
|
| 23 |
+
"source": [
|
| 24 |
+
"## 1. Install dependencies\n",
|
| 25 |
+
"\n",
|
| 26 |
+
"Unsloth pulls a custom torch + vllm + xformers stack tuned for free Colab T4. Pin trl to a version compatible with `GRPOTrainer` (>=0.14)."
|
| 27 |
+
]
|
| 28 |
+
},
|
| 29 |
+
{
|
| 30 |
+
"cell_type": "code",
|
| 31 |
+
"execution_count": null,
|
| 32 |
+
"metadata": {},
|
| 33 |
+
"outputs": [],
|
| 34 |
+
"source": [
|
| 35 |
+
"%%capture\n",
|
| 36 |
+
"!pip install --upgrade pip\n",
|
| 37 |
+
"!pip install unsloth vllm\n",
|
| 38 |
+
"!pip install --upgrade --no-deps \"trl>=0.14\" peft accelerate bitsandbytes\n",
|
| 39 |
+
"!pip install pydantic openenv huggingface_hub matplotlib"
|
| 40 |
+
]
|
| 41 |
+
},
|
| 42 |
+
{
|
| 43 |
+
"cell_type": "markdown",
|
| 44 |
+
"metadata": {},
|
| 45 |
+
"source": [
|
| 46 |
+
"## 2. Authenticate with Hugging Face\n",
|
| 47 |
+
"\n",
|
| 48 |
+
"Reads `HF_TOKEN` from Colab Secrets. Falls back to interactive login if not set."
|
| 49 |
+
]
|
| 50 |
+
},
|
| 51 |
+
{
|
| 52 |
+
"cell_type": "code",
|
| 53 |
+
"execution_count": null,
|
| 54 |
+
"metadata": {},
|
| 55 |
+
"outputs": [],
|
| 56 |
+
"source": [
|
| 57 |
+
"import os\n",
|
| 58 |
+
"\n",
|
| 59 |
+
"try:\n",
|
| 60 |
+
" from google.colab import userdata\n",
|
| 61 |
+
" HF_TOKEN = userdata.get(\"HF_TOKEN\")\n",
|
| 62 |
+
" os.environ[\"HF_TOKEN\"] = HF_TOKEN\n",
|
| 63 |
+
"except Exception:\n",
|
| 64 |
+
" from huggingface_hub import login\n",
|
| 65 |
+
" login()\n",
|
| 66 |
+
" HF_TOKEN = os.environ.get(\"HF_TOKEN\", \"\")\n",
|
| 67 |
+
"\n",
|
| 68 |
+
"assert HF_TOKEN, \"HF_TOKEN is required (set in Colab Secrets or via login())\"\n",
|
| 69 |
+
"print(\"HF auth OK\")"
|
| 70 |
+
]
|
| 71 |
+
},
|
| 72 |
+
{
|
| 73 |
+
"cell_type": "markdown",
|
| 74 |
+
"metadata": {},
|
| 75 |
+
"source": [
|
| 76 |
+
"## 3. Clone CrisisWorldCortex and install\n",
|
| 77 |
+
"\n",
|
| 78 |
+
"Pulls the deployed HF Space's repo and installs locally. Provides the `CrisisworldcortexEnv` HTTP client + the `baselines.flat_agent` system prompt + parser used by B1."
|
| 79 |
+
]
|
| 80 |
+
},
|
| 81 |
+
{
|
| 82 |
+
"cell_type": "code",
|
| 83 |
+
"execution_count": null,
|
| 84 |
+
"metadata": {},
|
| 85 |
+
"outputs": [],
|
| 86 |
+
"source": [
|
| 87 |
+
"%%capture\n",
|
| 88 |
+
"!rm -rf /content/CrisisWorldCortex\n",
|
| 89 |
+
"!git clone https://huggingface.co/spaces/Angshuman28/CrisisWorldCortex /content/CrisisWorldCortex\n",
|
| 90 |
+
"%cd /content/CrisisWorldCortex\n",
|
| 91 |
+
"!pip install -e ."
|
| 92 |
+
]
|
| 93 |
+
},
|
| 94 |
+
{
|
| 95 |
+
"cell_type": "code",
|
| 96 |
+
"execution_count": null,
|
| 97 |
+
"metadata": {},
|
| 98 |
+
"outputs": [],
|
| 99 |
+
"source": [
|
| 100 |
+
"# Sanity: imports resolve, env client constructs.\n",
|
| 101 |
+
"import sys\n",
|
| 102 |
+
"sys.path.insert(0, \"/content/CrisisWorldCortex\")\n",
|
| 103 |
+
"\n",
|
| 104 |
+
"from CrisisWorldCortex import CrisisworldcortexAction, CrisisworldcortexObservation\n",
|
| 105 |
+
"from CrisisWorldCortex.client import CrisisworldcortexEnv\n",
|
| 106 |
+
"from baselines.flat_agent import (\n",
|
| 107 |
+
" build_system_prompt,\n",
|
| 108 |
+
" parse_action,\n",
|
| 109 |
+
" parse_failure_marker,\n",
|
| 110 |
+
" serialize_observation,\n",
|
| 111 |
+
")\n",
|
| 112 |
+
"print(\"CrisisWorld imports OK\")"
|
| 113 |
+
]
|
| 114 |
+
},
|
| 115 |
+
{
|
| 116 |
+
"cell_type": "markdown",
|
| 117 |
+
"metadata": {},
|
| 118 |
+
"source": [
|
| 119 |
+
"## 4. Load Qwen3-1.7B with LoRA via Unsloth\n",
|
| 120 |
+
"\n",
|
| 121 |
+
"Qwen3-1.7B fits comfortably on a T4 with 4-bit quantization. LoRA rank 32 — enough to learn the JSON-action format and modest policy improvements; cheap to merge."
|
| 122 |
+
]
|
| 123 |
+
},
|
| 124 |
+
{
|
| 125 |
+
"cell_type": "code",
|
| 126 |
+
"execution_count": null,
|
| 127 |
+
"metadata": {},
|
| 128 |
+
"outputs": [],
|
| 129 |
+
"source": [
|
| 130 |
+
"from unsloth import FastLanguageModel\n",
|
| 131 |
+
"import torch\n",
|
| 132 |
+
"\n",
|
| 133 |
+
"MAX_SEQ_LEN = 4096\n",
|
| 134 |
+
"MODEL_NAME = \"unsloth/Qwen3-1.7B\"\n",
|
| 135 |
+
"\n",
|
| 136 |
+
"model, tokenizer = FastLanguageModel.from_pretrained(\n",
|
| 137 |
+
" model_name=MODEL_NAME,\n",
|
| 138 |
+
" max_seq_length=MAX_SEQ_LEN,\n",
|
| 139 |
+
" load_in_4bit=True,\n",
|
| 140 |
+
" fast_inference=True, # vLLM-backed generate, required by GRPOTrainer\n",
|
| 141 |
+
" max_lora_rank=32,\n",
|
| 142 |
+
" gpu_memory_utilization=0.6,\n",
|
| 143 |
+
")\n",
|
| 144 |
+
"\n",
|
| 145 |
+
"model = FastLanguageModel.get_peft_model(\n",
|
| 146 |
+
" model,\n",
|
| 147 |
+
" r=32,\n",
|
| 148 |
+
" target_modules=[\"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\",\n",
|
| 149 |
+
" \"gate_proj\", \"up_proj\", \"down_proj\"],\n",
|
| 150 |
+
" lora_alpha=64,\n",
|
| 151 |
+
" use_gradient_checkpointing=\"unsloth\",\n",
|
| 152 |
+
" random_state=42,\n",
|
| 153 |
+
")\n",
|
| 154 |
+
"print(\"Model + LoRA loaded\")"
|
| 155 |
+
]
|
| 156 |
+
},
|
| 157 |
+
{
|
| 158 |
+
"cell_type": "markdown",
|
| 159 |
+
"metadata": {},
|
| 160 |
+
"source": [
|
| 161 |
+
"## 5. Connect to the deployed CrisisWorld env\n",
|
| 162 |
+
"\n",
|
| 163 |
+
"Uses the public HF Space URL. Each rollout calls `env.reset()` then `env.step(action)` once."
|
| 164 |
+
]
|
| 165 |
+
},
|
| 166 |
+
{
|
| 167 |
+
"cell_type": "code",
|
| 168 |
+
"execution_count": null,
|
| 169 |
+
"metadata": {},
|
| 170 |
+
"outputs": [],
|
| 171 |
+
"source": [
|
| 172 |
+
"ENV_URL = \"https://angshuman28-crisisworldcortex.hf.space\"\n",
|
| 173 |
+
"TASKS = (\"outbreak_easy\", \"outbreak_medium\", \"outbreak_hard\")\n",
|
| 174 |
+
"EPISODE_TICKS = 12\n",
|
| 175 |
+
"\n",
|
| 176 |
+
"def make_env() -> CrisisworldcortexEnv:\n",
|
| 177 |
+
" return CrisisworldcortexEnv(base_url=ENV_URL)\n",
|
| 178 |
+
"\n",
|
| 179 |
+
"_test_env = make_env()\n",
|
| 180 |
+
"_obs = _test_env.reset(task_name=\"outbreak_easy\", seed=0, max_ticks=EPISODE_TICKS)\n",
|
| 181 |
+
"print(f\"Env OK. Initial tick={_obs.tick}, regions={[r.region for r in _obs.regions]}\")"
|
| 182 |
+
]
|
| 183 |
+
},
|
| 184 |
+
{
|
| 185 |
+
"cell_type": "markdown",
|
| 186 |
+
"metadata": {},
|
| 187 |
+
"source": [
|
| 188 |
+
"## 6. Build the prompt dataset and reward function\n",
|
| 189 |
+
"\n",
|
| 190 |
+
"Each example in the dataset is a `(task, seed)` pair. The reward function:\n",
|
| 191 |
+
"1. Resets the env to that `(task, seed)`.\n",
|
| 192 |
+
"2. Parses the model's completion as a `OuterActionPayload`.\n",
|
| 193 |
+
"3. Submits to the env, returns `obs.reward` (post-Phase-1 range `[-1, 1]`).\n",
|
| 194 |
+
"4. On parse failure, submits `parse_failure_marker()` so the §19 -1.0 + terminate contract fires."
|
| 195 |
+
]
|
| 196 |
+
},
|
| 197 |
+
{
|
| 198 |
+
"cell_type": "code",
|
| 199 |
+
"execution_count": null,
|
| 200 |
+
"metadata": {},
|
| 201 |
+
"outputs": [],
|
| 202 |
+
"source": [
|
| 203 |
+
"from datasets import Dataset\n",
|
| 204 |
+
"import random\n",
|
| 205 |
+
"\n",
|
| 206 |
+
"SYSTEM_PROMPT = build_system_prompt()\n",
|
| 207 |
+
"\n",
|
| 208 |
+
"def build_user_prompt(obs: CrisisworldcortexObservation) -> str:\n",
|
| 209 |
+
" return serialize_observation(obs)\n",
|
| 210 |
+
"\n",
|
| 211 |
+
"def make_chat_prompt(obs: CrisisworldcortexObservation) -> str:\n",
|
| 212 |
+
" return tokenizer.apply_chat_template(\n",
|
| 213 |
+
" [\n",
|
| 214 |
+
" {\"role\": \"system\", \"content\": SYSTEM_PROMPT},\n",
|
| 215 |
+
" {\"role\": \"user\", \"content\": build_user_prompt(obs)},\n",
|
| 216 |
+
" ],\n",
|
| 217 |
+
" tokenize=False,\n",
|
| 218 |
+
" add_generation_prompt=True,\n",
|
| 219 |
+
" )\n",
|
| 220 |
+
"\n",
|
| 221 |
+
"rng = random.Random(0)\n",
|
| 222 |
+
"_seed_pool = []\n",
|
| 223 |
+
"for task in TASKS:\n",
|
| 224 |
+
" for seed in range(50):\n",
|
| 225 |
+
" _seed_pool.append({\"task\": task, \"seed\": seed})\n",
|
| 226 |
+
"rng.shuffle(_seed_pool)\n",
|
| 227 |
+
"\n",
|
| 228 |
+
"_prompts = []\n",
|
| 229 |
+
"_meta = []\n",
|
| 230 |
+
"for entry in _seed_pool:\n",
|
| 231 |
+
" env = make_env()\n",
|
| 232 |
+
" obs = env.reset(task_name=entry[\"task\"], seed=entry[\"seed\"], max_ticks=EPISODE_TICKS)\n",
|
| 233 |
+
" _prompts.append(make_chat_prompt(obs))\n",
|
| 234 |
+
" _meta.append(entry)\n",
|
| 235 |
+
"\n",
|
| 236 |
+
"train_dataset = Dataset.from_dict({\n",
|
| 237 |
+
" \"prompt\": _prompts,\n",
|
| 238 |
+
" \"task\": [m[\"task\"] for m in _meta],\n",
|
| 239 |
+
" \"seed\": [m[\"seed\"] for m in _meta],\n",
|
| 240 |
+
"})\n",
|
| 241 |
+
"print(f\"Dataset built: {len(train_dataset)} examples\")"
|
| 242 |
+
]
|
| 243 |
+
},
|
| 244 |
+
{
|
| 245 |
+
"cell_type": "code",
|
| 246 |
+
"execution_count": null,
|
| 247 |
+
"metadata": {},
|
| 248 |
+
"outputs": [],
|
| 249 |
+
"source": [
|
| 250 |
+
"def crisisworld_reward(\n",
|
| 251 |
+
" prompts: list[str],\n",
|
| 252 |
+
" completions: list[str],\n",
|
| 253 |
+
" task: list[str],\n",
|
| 254 |
+
" seed: list[int],\n",
|
| 255 |
+
" **_kwargs,\n",
|
| 256 |
+
") -> list[float]:\n",
|
| 257 |
+
" \"\"\"GRPO reward function: one env step per (prompt, completion) pair.\n",
|
| 258 |
+
"\n",
|
| 259 |
+
" Reward source: Phase-1-fixed env outer_reward in [-1, 1].\n",
|
| 260 |
+
" Parse failure → submits parse_failure_marker → r_policy = -1.0 + terminate.\n",
|
| 261 |
+
" \"\"\"\n",
|
| 262 |
+
" rewards: list[float] = []\n",
|
| 263 |
+
" for completion, t, s in zip(completions, task, seed):\n",
|
| 264 |
+
" env = make_env()\n",
|
| 265 |
+
" env.reset(task_name=t, seed=int(s), max_ticks=EPISODE_TICKS)\n",
|
| 266 |
+
" action_payload = parse_action(completion)\n",
|
| 267 |
+
" if action_payload is None:\n",
|
| 268 |
+
" action_payload = parse_failure_marker()\n",
|
| 269 |
+
" try:\n",
|
| 270 |
+
" result = env.step(CrisisworldcortexAction(action=action_payload))\n",
|
| 271 |
+
" reward = result.observation.reward if hasattr(result, \"observation\") else result.reward\n",
|
| 272 |
+
" rewards.append(float(reward) if reward is not None else 0.0)\n",
|
| 273 |
+
" except Exception as exc:\n",
|
| 274 |
+
" print(f\"[WARN] env.step failed task={t} seed={s}: {exc}\")\n",
|
| 275 |
+
" rewards.append(-1.0)\n",
|
| 276 |
+
" return rewards"
|
| 277 |
+
]
|
| 278 |
+
},
|
| 279 |
+
{
|
| 280 |
+
"cell_type": "markdown",
|
| 281 |
+
"metadata": {},
|
| 282 |
+
"source": [
|
| 283 |
+
"## 7. GRPO training\n",
|
| 284 |
+
"\n",
|
| 285 |
+
"300 steps × group size 4 = 1200 rollouts. Each rollout is one HF Space round-trip (~1s) — total ~20–30 min on T4."
|
| 286 |
+
]
|
| 287 |
+
},
|
| 288 |
+
{
|
| 289 |
+
"cell_type": "code",
|
| 290 |
+
"execution_count": null,
|
| 291 |
+
"metadata": {},
|
| 292 |
+
"outputs": [],
|
| 293 |
+
"source": [
|
| 294 |
+
"from trl import GRPOConfig, GRPOTrainer\n",
|
| 295 |
+
"\n",
|
| 296 |
+
"MAX_TRAIN_STEPS = 300\n",
|
| 297 |
+
"GROUP_SIZE = 4\n",
|
| 298 |
+
"MAX_PROMPT_LEN = 2048\n",
|
| 299 |
+
"MAX_COMPLETION_LEN = 512\n",
|
| 300 |
+
"\n",
|
| 301 |
+
"training_args = GRPOConfig(\n",
|
| 302 |
+
" output_dir=\"/content/b1_grpo_output\",\n",
|
| 303 |
+
" learning_rate=5e-6,\n",
|
| 304 |
+
" per_device_train_batch_size=GROUP_SIZE,\n",
|
| 305 |
+
" gradient_accumulation_steps=1,\n",
|
| 306 |
+
" num_generations=GROUP_SIZE,\n",
|
| 307 |
+
" max_prompt_length=MAX_PROMPT_LEN,\n",
|
| 308 |
+
" max_completion_length=MAX_COMPLETION_LEN,\n",
|
| 309 |
+
" max_steps=MAX_TRAIN_STEPS,\n",
|
| 310 |
+
" save_steps=100,\n",
|
| 311 |
+
" logging_steps=5,\n",
|
| 312 |
+
" report_to=\"none\",\n",
|
| 313 |
+
" bf16=True,\n",
|
| 314 |
+
" optim=\"adamw_8bit\",\n",
|
| 315 |
+
" temperature=0.8,\n",
|
| 316 |
+
" use_vllm=True,\n",
|
| 317 |
+
" vllm_mode=\"colocate\",\n",
|
| 318 |
+
" seed=42,\n",
|
| 319 |
+
")\n",
|
| 320 |
+
"\n",
|
| 321 |
+
"trainer = GRPOTrainer(\n",
|
| 322 |
+
" model=model,\n",
|
| 323 |
+
" processing_class=tokenizer,\n",
|
| 324 |
+
" reward_funcs=[crisisworld_reward],\n",
|
| 325 |
+
" args=training_args,\n",
|
| 326 |
+
" train_dataset=train_dataset,\n",
|
| 327 |
+
")\n",
|
| 328 |
+
"print(\"GRPOTrainer constructed; starting train()...\")"
|
| 329 |
+
]
|
| 330 |
+
},
|
| 331 |
+
{
|
| 332 |
+
"cell_type": "code",
|
| 333 |
+
"execution_count": null,
|
| 334 |
+
"metadata": {},
|
| 335 |
+
"outputs": [],
|
| 336 |
+
"source": [
|
| 337 |
+
"trainer.train()"
|
| 338 |
+
]
|
| 339 |
+
},
|
| 340 |
+
{
|
| 341 |
+
"cell_type": "markdown",
|
| 342 |
+
"metadata": {},
|
| 343 |
+
"source": [
|
| 344 |
+
"## 8. Save the trained LoRA adapter to HF Hub\n",
|
| 345 |
+
"\n",
|
| 346 |
+
"Pushes to `<your-username>/crisisworld-b1-grpo-qwen3-1p7b`. Change the repo name below if you want a different namespace."
|
| 347 |
+
]
|
| 348 |
+
},
|
| 349 |
+
{
|
| 350 |
+
"cell_type": "code",
|
| 351 |
+
"execution_count": null,
|
| 352 |
+
"metadata": {},
|
| 353 |
+
"outputs": [],
|
| 354 |
+
"source": [
|
| 355 |
+
"from huggingface_hub import HfApi\n",
|
| 356 |
+
"\n",
|
| 357 |
+
"HUB_REPO = \"Angshuman28/crisisworld-b1-grpo-qwen3-1p7b\"\n",
|
| 358 |
+
"\n",
|
| 359 |
+
"model.save_pretrained(\"/content/b1_grpo_lora\")\n",
|
| 360 |
+
"tokenizer.save_pretrained(\"/content/b1_grpo_lora\")\n",
|
| 361 |
+
"\n",
|
| 362 |
+
"api = HfApi()\n",
|
| 363 |
+
"api.create_repo(HUB_REPO, exist_ok=True, repo_type=\"model\", private=False, token=HF_TOKEN)\n",
|
| 364 |
+
"api.upload_folder(\n",
|
| 365 |
+
" folder_path=\"/content/b1_grpo_lora\",\n",
|
| 366 |
+
" repo_id=HUB_REPO,\n",
|
| 367 |
+
" repo_type=\"model\",\n",
|
| 368 |
+
" token=HF_TOKEN,\n",
|
| 369 |
+
")\n",
|
| 370 |
+
"print(f\"Saved to https://huggingface.co/{HUB_REPO}\")"
|
| 371 |
+
]
|
| 372 |
+
},
|
| 373 |
+
{
|
| 374 |
+
"cell_type": "markdown",
|
| 375 |
+
"metadata": {},
|
| 376 |
+
"source": [
|
| 377 |
+
"## 9. Eval: trained adapter vs base model on 3 tasks\n",
|
| 378 |
+
"\n",
|
| 379 |
+
"Runs a single full episode (12 ticks) per task, per model. Reports cumulative reward."
|
| 380 |
+
]
|
| 381 |
+
},
|
| 382 |
+
{
|
| 383 |
+
"cell_type": "code",
|
| 384 |
+
"execution_count": null,
|
| 385 |
+
"metadata": {},
|
| 386 |
+
"outputs": [],
|
| 387 |
+
"source": [
|
| 388 |
+
"def _hf_chat(model_inst, tokenizer_inst, system: str, user: str, max_new_tokens: int = 256) -> str:\n",
|
| 389 |
+
" prompt = tokenizer_inst.apply_chat_template(\n",
|
| 390 |
+
" [{\"role\": \"system\", \"content\": system}, {\"role\": \"user\", \"content\": user}],\n",
|
| 391 |
+
" tokenize=False,\n",
|
| 392 |
+
" add_generation_prompt=True,\n",
|
| 393 |
+
" )\n",
|
| 394 |
+
" inputs = tokenizer_inst(prompt, return_tensors=\"pt\").to(model_inst.device)\n",
|
| 395 |
+
" with torch.no_grad():\n",
|
| 396 |
+
" out = model_inst.generate(\n",
|
| 397 |
+
" **inputs, max_new_tokens=max_new_tokens, do_sample=False, temperature=0.0,\n",
|
| 398 |
+
" )\n",
|
| 399 |
+
" return tokenizer_inst.decode(out[0][inputs[\"input_ids\"].shape[1]:], skip_special_tokens=True)\n",
|
| 400 |
+
"\n",
|
| 401 |
+
"def run_one_episode(model_inst, tokenizer_inst, task: str, seed: int) -> float:\n",
|
| 402 |
+
" env = make_env()\n",
|
| 403 |
+
" obs = env.reset(task_name=task, seed=seed, max_ticks=EPISODE_TICKS)\n",
|
| 404 |
+
" cumulative = 0.0\n",
|
| 405 |
+
" for tick in range(EPISODE_TICKS):\n",
|
| 406 |
+
" completion = _hf_chat(model_inst, tokenizer_inst, SYSTEM_PROMPT, serialize_observation(obs))\n",
|
| 407 |
+
" action = parse_action(completion) or parse_failure_marker()\n",
|
| 408 |
+
" result = env.step(CrisisworldcortexAction(action=action))\n",
|
| 409 |
+
" obs = result.observation if hasattr(result, \"observation\") else result\n",
|
| 410 |
+
" reward = obs.reward if obs.reward is not None else 0.0\n",
|
| 411 |
+
" cumulative += reward\n",
|
| 412 |
+
" if obs.done:\n",
|
| 413 |
+
" break\n",
|
| 414 |
+
" return cumulative\n",
|
| 415 |
+
"\n",
|
| 416 |
+
"FastLanguageModel.for_inference(model)\n",
|
| 417 |
+
"trained_results = {t: run_one_episode(model, tokenizer, t, seed=0) for t in TASKS}\n",
|
| 418 |
+
"print(f\"Trained model cumulative reward per task: {trained_results}\")"
|
| 419 |
+
]
|
| 420 |
+
},
|
| 421 |
+
{
|
| 422 |
+
"cell_type": "code",
|
| 423 |
+
"execution_count": null,
|
| 424 |
+
"metadata": {},
|
| 425 |
+
"outputs": [],
|
| 426 |
+
"source": [
|
| 427 |
+
"# Reload base Qwen3-1.7B (no LoRA) for the comparison.\n",
|
| 428 |
+
"base_model, base_tokenizer = FastLanguageModel.from_pretrained(\n",
|
| 429 |
+
" model_name=MODEL_NAME,\n",
|
| 430 |
+
" max_seq_length=MAX_SEQ_LEN,\n",
|
| 431 |
+
" load_in_4bit=True,\n",
|
| 432 |
+
" fast_inference=False,\n",
|
| 433 |
+
")\n",
|
| 434 |
+
"FastLanguageModel.for_inference(base_model)\n",
|
| 435 |
+
"base_results = {t: run_one_episode(base_model, base_tokenizer, t, seed=0) for t in TASKS}\n",
|
| 436 |
+
"print(f\"Base model cumulative reward per task: {base_results}\")"
|
| 437 |
+
]
|
| 438 |
+
},
|
| 439 |
+
{
|
| 440 |
+
"cell_type": "markdown",
|
| 441 |
+
"metadata": {},
|
| 442 |
+
"source": [
|
| 443 |
+
"## 10. Plot eval comparison\n",
|
| 444 |
+
"\n",
|
| 445 |
+
"Bar chart: trained vs base, cumulative episode reward by task."
|
| 446 |
+
]
|
| 447 |
+
},
|
| 448 |
+
{
|
| 449 |
+
"cell_type": "code",
|
| 450 |
+
"execution_count": null,
|
| 451 |
+
"metadata": {},
|
| 452 |
+
"outputs": [],
|
| 453 |
+
"source": [
|
| 454 |
+
"import matplotlib.pyplot as plt\n",
|
| 455 |
+
"import numpy as np\n",
|
| 456 |
+
"\n",
|
| 457 |
+
"task_names = list(TASKS)\n",
|
| 458 |
+
"trained_vals = [trained_results[t] for t in task_names]\n",
|
| 459 |
+
"base_vals = [base_results[t] for t in task_names]\n",
|
| 460 |
+
"\n",
|
| 461 |
+
"x = np.arange(len(task_names))\n",
|
| 462 |
+
"width = 0.35\n",
|
| 463 |
+
"\n",
|
| 464 |
+
"fig, ax = plt.subplots(figsize=(9, 5))\n",
|
| 465 |
+
"ax.bar(x - width/2, base_vals, width, label=\"Base Qwen3-1.7B\")\n",
|
| 466 |
+
"ax.bar(x + width/2, trained_vals, width, label=\"GRPO-trained Qwen3-1.7B\")\n",
|
| 467 |
+
"ax.set_xticks(x)\n",
|
| 468 |
+
"ax.set_xticklabels(task_names)\n",
|
| 469 |
+
"ax.set_ylabel(\"Cumulative episode reward\")\n",
|
| 470 |
+
"ax.set_title(\"B1 GRPO: trained vs base\")\n",
|
| 471 |
+
"ax.legend()\n",
|
| 472 |
+
"ax.axhline(0.0, linestyle=\":\", color=\"grey\")\n",
|
| 473 |
+
"plt.tight_layout()\n",
|
| 474 |
+
"plt.show()"
|
| 475 |
+
]
|
| 476 |
+
}
|
| 477 |
+
],
|
| 478 |
+
"metadata": {
|
| 479 |
+
"kernelspec": {
|
| 480 |
+
"display_name": "Python 3",
|
| 481 |
+
"language": "python",
|
| 482 |
+
"name": "python3"
|
| 483 |
+
},
|
| 484 |
+
"language_info": {
|
| 485 |
+
"name": "python",
|
| 486 |
+
"version": "3.10"
|
| 487 |
+
},
|
| 488 |
+
"accelerator": "GPU",
|
| 489 |
+
"colab": {
|
| 490 |
+
"gpuType": "T4",
|
| 491 |
+
"provenance": []
|
| 492 |
+
}
|
| 493 |
+
},
|
| 494 |
+
"nbformat": 4,
|
| 495 |
+
"nbformat_minor": 5
|
| 496 |
+
}
|
notebooks/train_cortex_router_grpo.ipynb
ADDED
|
@@ -0,0 +1,567 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"cells": [
|
| 3 |
+
{
|
| 4 |
+
"cell_type": "markdown",
|
| 5 |
+
"id": "7fb27b941602401d91542211134fc71a",
|
| 6 |
+
"metadata": {},
|
| 7 |
+
"source": [
|
| 8 |
+
"# Cortex Small-LLM Router GRPO training (Workstream B Phase 4 — SKELETON, post-pivot)\n",
|
| 9 |
+
"\n",
|
| 10 |
+
"**Status:** scaffold with explicit `TODO(conv-point)` markers at every Cortex-dependent integration site. **Do not run end-to-end yet** — the cells marked TODO will fail until Workstream A's Session 13 ships:\n",
|
| 11 |
+
"\n",
|
| 12 |
+
"- `cortex.metacognition.MetacognitionState` — featurization source.\n",
|
| 13 |
+
"- `cortex.routing_policy.RoutingPolicy` — trainable interface (Phase A §6).\n",
|
| 14 |
+
"- `cortex.council.Council` — deliberation orchestrator that drives rollouts.\n",
|
| 15 |
+
"- `baselines.cortex_fixed_router.B3CortexFixedRouter` — generates the deterministic-router corpus.\n",
|
| 16 |
+
"\n",
|
| 17 |
+
"**Architecture (post-pivot — replaces the MLP-head approach in commit `5489e55`):**\n",
|
| 18 |
+
"- Router is a small LLM: `unsloth/Qwen3-1.5B-Instruct` + LoRA rank 16. Roughly 3 GB on a100-large in 4-bit.\n",
|
| 19 |
+
"- Input: NL summary of `MetacognitionState` (~300 tokens).\n",
|
| 20 |
+
"- Output: structured JSON matching `cortex.schemas.RoutingAction` (~150 tokens).\n",
|
| 21 |
+
"- GRPO via the same Unsloth + TRL `GRPOTrainer` pipeline as the B1 notebook (one less code path).\n",
|
| 22 |
+
"- Reward: Phase-1-fixed `outer_reward` ∈ [-1, 1] composed with the token-budget penalty via `training.reward_shaping.shape_reward`.\n",
|
| 23 |
+
"\n",
|
| 24 |
+
"**Why a small LLM (not an MLP)?** With A100 compute available, a small LLM gives better reasoning over complex MetacognitionState, is interpretable in the pitch demo (you can read what the router thinks), and reuses the exact training infrastructure as the B1 baseline."
|
| 25 |
+
]
|
| 26 |
+
},
|
| 27 |
+
{
|
| 28 |
+
"cell_type": "markdown",
|
| 29 |
+
"id": "acae54e37e7d407bbb7b55eff062a284",
|
| 30 |
+
"metadata": {},
|
| 31 |
+
"source": [
|
| 32 |
+
"## 1. Install dependencies\n",
|
| 33 |
+
"\n",
|
| 34 |
+
"Same Unsloth + vLLM + TRL stack as the B1 notebook; smaller GPU footprint because the trainable router is 1.5B."
|
| 35 |
+
]
|
| 36 |
+
},
|
| 37 |
+
{
|
| 38 |
+
"cell_type": "code",
|
| 39 |
+
"execution_count": null,
|
| 40 |
+
"id": "9a63283cbaf04dbcab1f6479b197f3a8",
|
| 41 |
+
"metadata": {},
|
| 42 |
+
"outputs": [],
|
| 43 |
+
"source": [
|
| 44 |
+
"%%capture\n",
|
| 45 |
+
"!pip install --upgrade pip\n",
|
| 46 |
+
"!pip install unsloth vllm\n",
|
| 47 |
+
"!pip install --upgrade --no-deps \"trl>=0.14\" peft accelerate bitsandbytes\n",
|
| 48 |
+
"!pip install pydantic openenv huggingface_hub matplotlib"
|
| 49 |
+
]
|
| 50 |
+
},
|
| 51 |
+
{
|
| 52 |
+
"cell_type": "markdown",
|
| 53 |
+
"id": "8dd0d8092fe74a7c96281538738b07e2",
|
| 54 |
+
"metadata": {},
|
| 55 |
+
"source": [
|
| 56 |
+
"## 2. Authenticate with Hugging Face"
|
| 57 |
+
]
|
| 58 |
+
},
|
| 59 |
+
{
|
| 60 |
+
"cell_type": "code",
|
| 61 |
+
"execution_count": null,
|
| 62 |
+
"id": "72eea5119410473aa328ad9291626812",
|
| 63 |
+
"metadata": {},
|
| 64 |
+
"outputs": [],
|
| 65 |
+
"source": [
|
| 66 |
+
"import os\n",
|
| 67 |
+
"\n",
|
| 68 |
+
"try:\n",
|
| 69 |
+
" from google.colab import userdata\n",
|
| 70 |
+
"\n",
|
| 71 |
+
" HF_TOKEN = userdata.get(\"HF_TOKEN\")\n",
|
| 72 |
+
" os.environ[\"HF_TOKEN\"] = HF_TOKEN\n",
|
| 73 |
+
"except Exception:\n",
|
| 74 |
+
" from huggingface_hub import login\n",
|
| 75 |
+
"\n",
|
| 76 |
+
" login()\n",
|
| 77 |
+
" HF_TOKEN = os.environ.get(\"HF_TOKEN\", \"\")\n",
|
| 78 |
+
"\n",
|
| 79 |
+
"assert HF_TOKEN, \"HF_TOKEN required\""
|
| 80 |
+
]
|
| 81 |
+
},
|
| 82 |
+
{
|
| 83 |
+
"cell_type": "markdown",
|
| 84 |
+
"id": "8edb47106e1a46a883d545849b8ab81b",
|
| 85 |
+
"metadata": {},
|
| 86 |
+
"source": [
|
| 87 |
+
"## 3. Clone CrisisWorldCortex (post-Cortex-Session-13 deploy)\n",
|
| 88 |
+
"\n",
|
| 89 |
+
"**TODO(conv-point):** the target Space repo at convergence point will have `cortex/*` populated through Session 13 (subagents + lenses + brains + council + metacognition + routing_policy + B3). Until then, the Cortex-dependent imports below will fail."
|
| 90 |
+
]
|
| 91 |
+
},
|
| 92 |
+
{
|
| 93 |
+
"cell_type": "code",
|
| 94 |
+
"execution_count": null,
|
| 95 |
+
"id": "10185d26023b46108eb7d9f57d49d2b3",
|
| 96 |
+
"metadata": {},
|
| 97 |
+
"outputs": [],
|
| 98 |
+
"source": [
|
| 99 |
+
"%%capture\n",
|
| 100 |
+
"!rm -rf /content/CrisisWorldCortex\n",
|
| 101 |
+
"!git clone https://huggingface.co/spaces/Angshuman28/CrisisWorldCortex /content/CrisisWorldCortex\n",
|
| 102 |
+
"%cd /content/CrisisWorldCortex\n",
|
| 103 |
+
"!pip install -e ."
|
| 104 |
+
]
|
| 105 |
+
},
|
| 106 |
+
{
|
| 107 |
+
"cell_type": "code",
|
| 108 |
+
"execution_count": null,
|
| 109 |
+
"id": "8763a12b2bbd4a93a75aff182afb95dc",
|
| 110 |
+
"metadata": {},
|
| 111 |
+
"outputs": [],
|
| 112 |
+
"source": [
|
| 113 |
+
"import sys\n",
|
| 114 |
+
"\n",
|
| 115 |
+
"sys.path.insert(0, \"/content/CrisisWorldCortex\")\n",
|
| 116 |
+
"\n",
|
| 117 |
+
"\n",
|
| 118 |
+
"# TODO(conv-point): uncomment when Cortex Session 13 lands.\n",
|
| 119 |
+
"# from cortex.schemas import MetacognitionState, RoutingAction, RouterStep\n",
|
| 120 |
+
"# from cortex.routing_policy import RoutingPolicy\n",
|
| 121 |
+
"# from cortex.council import Council\n",
|
| 122 |
+
"# from baselines.cortex_fixed_router import B3CortexFixedRouter\n",
|
| 123 |
+
"print(\"Phase-2 training utilities OK; Cortex imports gated until Session 13\")"
|
| 124 |
+
]
|
| 125 |
+
},
|
| 126 |
+
{
|
| 127 |
+
"cell_type": "markdown",
|
| 128 |
+
"id": "7623eae2785240b9bd12b16a66d81610",
|
| 129 |
+
"metadata": {},
|
| 130 |
+
"source": [
|
| 131 |
+
"## 4. Load Qwen3-1.5B-Instruct (router) with LoRA\n",
|
| 132 |
+
"\n",
|
| 133 |
+
"Small enough to fit alongside frozen 7B/8B brain LLMs on a100-large (80GB). LoRA rank 16 — tighter than the B1 notebook's 32 because the action space is structured JSON (small effective vocabulary)."
|
| 134 |
+
]
|
| 135 |
+
},
|
| 136 |
+
{
|
| 137 |
+
"cell_type": "code",
|
| 138 |
+
"execution_count": null,
|
| 139 |
+
"id": "7cdc8c89c7104fffa095e18ddfef8986",
|
| 140 |
+
"metadata": {},
|
| 141 |
+
"outputs": [],
|
| 142 |
+
"source": [
|
| 143 |
+
"from unsloth import FastLanguageModel\n",
|
| 144 |
+
"\n",
|
| 145 |
+
"ROUTER_MODEL = \"unsloth/Qwen3-1.5B-Instruct-bnb-4bit\"\n",
|
| 146 |
+
"MAX_SEQ_LEN = 2048\n",
|
| 147 |
+
"\n",
|
| 148 |
+
"router_model, router_tokenizer = FastLanguageModel.from_pretrained(\n",
|
| 149 |
+
" model_name=ROUTER_MODEL,\n",
|
| 150 |
+
" max_seq_length=MAX_SEQ_LEN,\n",
|
| 151 |
+
" load_in_4bit=True,\n",
|
| 152 |
+
" fast_inference=True,\n",
|
| 153 |
+
" max_lora_rank=16,\n",
|
| 154 |
+
" gpu_memory_utilization=0.5, # share GPU with frozen brain LLMs at conv-point\n",
|
| 155 |
+
")\n",
|
| 156 |
+
"\n",
|
| 157 |
+
"router_model = FastLanguageModel.get_peft_model(\n",
|
| 158 |
+
" router_model,\n",
|
| 159 |
+
" r=16,\n",
|
| 160 |
+
" target_modules=[\"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\", \"gate_proj\", \"up_proj\", \"down_proj\"],\n",
|
| 161 |
+
" lora_alpha=32,\n",
|
| 162 |
+
" use_gradient_checkpointing=\"unsloth\",\n",
|
| 163 |
+
" random_state=42,\n",
|
| 164 |
+
")\n",
|
| 165 |
+
"print(\"Router (Qwen3-1.5B + LoRA r=16) loaded\")"
|
| 166 |
+
]
|
| 167 |
+
},
|
| 168 |
+
{
|
| 169 |
+
"cell_type": "markdown",
|
| 170 |
+
"id": "b118ea5561624da68c537baed56e602f",
|
| 171 |
+
"metadata": {},
|
| 172 |
+
"source": [
|
| 173 |
+
"## 5. Featurization: MetacognitionState → NL prompt + RoutingAction schema\n",
|
| 174 |
+
"\n",
|
| 175 |
+
"Replaces the MLP version's `(24,) np.float32` featurization. The router consumes a natural-language summary; output is structured JSON validated against `cortex.schemas.RoutingAction`.\n",
|
| 176 |
+
"\n",
|
| 177 |
+
"**TODO(conv-point):** the schema below references `cortex.schemas.MetacognitionState` (Session 13).\n",
|
| 178 |
+
"Per Phase A `cortex/schemas.py` and `docs/CORTEX_ARCHITECTURE.md` §6, the NL summary covers the 11 documented fields + the phase string."
|
| 179 |
+
]
|
| 180 |
+
},
|
| 181 |
+
{
|
| 182 |
+
"cell_type": "code",
|
| 183 |
+
"execution_count": null,
|
| 184 |
+
"id": "938c804e27f84196a10c8828c723f798",
|
| 185 |
+
"metadata": {},
|
| 186 |
+
"outputs": [],
|
| 187 |
+
"source": [
|
| 188 |
+
"ROUTER_SYSTEM_PROMPT = \"\"\"You are the Cortex router for CrisisWorldCortex. You receive a metacognition\n",
|
| 189 |
+
"state summary describing the current deliberation phase and emit ONE routing action as JSON.\n",
|
| 190 |
+
"\n",
|
| 191 |
+
"Allowed action kinds (per cortex/CLAUDE.md):\n",
|
| 192 |
+
"- call_subagent: invoke a brain's subagent. Required: brain (epidemiology|logistics|governance),\n",
|
| 193 |
+
" subagent (world_modeler|planner|critic).\n",
|
| 194 |
+
"- request_challenge: cross-brain critique. Required: challenger_brain, target_brain.\n",
|
| 195 |
+
"- switch_phase: advance the phase machine. Required: new_phase (divergence|challenge|narrowing|convergence).\n",
|
| 196 |
+
"- preserve_dissent: tag a minority recommendation. Required: tag (string, max 80 chars).\n",
|
| 197 |
+
"- emit_outer_action: close the tick with a final action.\n",
|
| 198 |
+
"- stop_and_no_op: close the tick with a no-op.\n",
|
| 199 |
+
"\n",
|
| 200 |
+
"Hard caps (binding): ≤2 deliberation rounds/tick, ≤1 cross-brain challenge/tick,\n",
|
| 201 |
+
"≤1 critic call per brain/tick, ≤6000 token budget per tick.\n",
|
| 202 |
+
"\n",
|
| 203 |
+
"Output exactly one JSON object — no markdown fences, no prose around it.\"\"\"\n",
|
| 204 |
+
"\n",
|
| 205 |
+
"PHASE_NAMES = (\"divergence\", \"challenge\", \"narrowing\", \"convergence\")\n",
|
| 206 |
+
"\n",
|
| 207 |
+
"\n",
|
| 208 |
+
"def metacog_state_to_prompt(state) -> str:\n",
|
| 209 |
+
" \"\"\"Convert MetacognitionState → NL summary for the router.\n",
|
| 210 |
+
"\n",
|
| 211 |
+
" TODO(conv-point): change ``state`` typing to MetacognitionState\n",
|
| 212 |
+
" (cortex.schemas) once Session 13 lands. Duck-typed for now.\n",
|
| 213 |
+
" \"\"\"\n",
|
| 214 |
+
" return (\n",
|
| 215 |
+
" f\"Tick {getattr(state, 'tick', 0)}, round {getattr(state, 'round', 1)}, \"\n",
|
| 216 |
+
" f\"phase={getattr(state, 'phase', 'divergence')}.\\n\"\n",
|
| 217 |
+
" f\"Inter-brain agreement: {getattr(state, 'inter_brain_agreement', 0.0):.2f}.\\n\"\n",
|
| 218 |
+
" f\"Average confidence: {getattr(state, 'average_confidence', 0.0):.2f}.\\n\"\n",
|
| 219 |
+
" f\"Average evidence support: {getattr(state, 'average_evidence_support', 0.0):.2f}.\\n\"\n",
|
| 220 |
+
" f\"Novelty yield (last round): {getattr(state, 'novelty_yield_last_round', 0.0):.2f}.\\n\"\n",
|
| 221 |
+
" f\"Collapse suspicion: {getattr(state, 'collapse_suspicion', 0.0):.2f}.\\n\"\n",
|
| 222 |
+
" f\"Budget remaining: {getattr(state, 'budget_remaining_frac', 1.0):.0%}.\\n\"\n",
|
| 223 |
+
" f\"Urgency: {getattr(state, 'urgency', 0.0):.2f}.\\n\"\n",
|
| 224 |
+
" f\"Preserved dissent count: {getattr(state, 'preserved_dissent_count', 0)}.\\n\"\n",
|
| 225 |
+
" f\"Cross-brain challenge used this tick: \"\n",
|
| 226 |
+
" f\"{bool(getattr(state, 'challenge_used_this_tick', 0))}.\\n\\n\"\n",
|
| 227 |
+
" f\"Choose the next routing action.\"\n",
|
| 228 |
+
" )\n",
|
| 229 |
+
"\n",
|
| 230 |
+
"\n",
|
| 231 |
+
"def make_router_chat_prompt(state) -> str:\n",
|
| 232 |
+
" return router_tokenizer.apply_chat_template(\n",
|
| 233 |
+
" [\n",
|
| 234 |
+
" {\"role\": \"system\", \"content\": ROUTER_SYSTEM_PROMPT},\n",
|
| 235 |
+
" {\"role\": \"user\", \"content\": metacog_state_to_prompt(state)},\n",
|
| 236 |
+
" ],\n",
|
| 237 |
+
" tokenize=False,\n",
|
| 238 |
+
" add_generation_prompt=True,\n",
|
| 239 |
+
" )"
|
| 240 |
+
]
|
| 241 |
+
},
|
| 242 |
+
{
|
| 243 |
+
"cell_type": "markdown",
|
| 244 |
+
"id": "504fb2a444614c0babb325280ed9130a",
|
| 245 |
+
"metadata": {},
|
| 246 |
+
"source": [
|
| 247 |
+
"## 6. Connect to the deployed CrisisWorld env\n",
|
| 248 |
+
"\n",
|
| 249 |
+
"Same env client pattern as the B1 notebook."
|
| 250 |
+
]
|
| 251 |
+
},
|
| 252 |
+
{
|
| 253 |
+
"cell_type": "code",
|
| 254 |
+
"execution_count": null,
|
| 255 |
+
"id": "59bbdb311c014d738909a11f9e486628",
|
| 256 |
+
"metadata": {},
|
| 257 |
+
"outputs": [],
|
| 258 |
+
"source": [
|
| 259 |
+
"from CrisisWorldCortex.client import CrisisworldcortexEnv\n",
|
| 260 |
+
"\n",
|
| 261 |
+
"ENV_URL = \"https://angshuman28-crisisworldcortex.hf.space\"\n",
|
| 262 |
+
"TASKS = (\"outbreak_easy\", \"outbreak_medium\", \"outbreak_hard\")\n",
|
| 263 |
+
"EPISODE_TICKS = 12\n",
|
| 264 |
+
"\n",
|
| 265 |
+
"\n",
|
| 266 |
+
"def make_env() -> CrisisworldcortexEnv:\n",
|
| 267 |
+
" return CrisisworldcortexEnv(base_url=ENV_URL)\n",
|
| 268 |
+
"\n",
|
| 269 |
+
"\n",
|
| 270 |
+
"_test_env = make_env()\n",
|
| 271 |
+
"_obs = _test_env.reset(task_name=\"outbreak_easy\", seed=0, max_ticks=EPISODE_TICKS)\n",
|
| 272 |
+
"print(f\"Env OK. Initial tick={_obs.tick}, regions={[r.region for r in _obs.regions]}\")"
|
| 273 |
+
]
|
| 274 |
+
},
|
| 275 |
+
{
|
| 276 |
+
"cell_type": "markdown",
|
| 277 |
+
"id": "b43b363d81ae4b689946ece5c682cd59",
|
| 278 |
+
"metadata": {},
|
| 279 |
+
"source": [
|
| 280 |
+
"## 7. Build training-data prompt set from B3 deterministic-router trajectories\n",
|
| 281 |
+
"\n",
|
| 282 |
+
"**TODO(conv-point):** B3CortexFixedRouter runs ~50 episodes. Each `RouterStep`'s `MetacognitionState` becomes a router-prompt; the GRPO completion is the router's emitted JSON.\n",
|
| 283 |
+
"\n",
|
| 284 |
+
"Replaces the MLP version's action-vocab + B3-trajectory-to-tuple pipeline. Same RolloutBuffer, different content shape."
|
| 285 |
+
]
|
| 286 |
+
},
|
| 287 |
+
{
|
| 288 |
+
"cell_type": "code",
|
| 289 |
+
"execution_count": null,
|
| 290 |
+
"id": "8a65eabff63a45729fe45fb5ade58bdc",
|
| 291 |
+
"metadata": {},
|
| 292 |
+
"outputs": [],
|
| 293 |
+
"source": [
|
| 294 |
+
"from datasets import Dataset\n",
|
| 295 |
+
"\n",
|
| 296 |
+
"TRAIN_EPISODES = 50\n",
|
| 297 |
+
"\n",
|
| 298 |
+
"\n",
|
| 299 |
+
"def collect_b3_router_prompts(num_episodes: int = TRAIN_EPISODES) -> Dataset:\n",
|
| 300 |
+
" \"\"\"Run B3 and convert each RouterStep into a (prompt, task, seed) row.\n",
|
| 301 |
+
"\n",
|
| 302 |
+
" TODO(conv-point): uncomment the body when B3CortexFixedRouter ships.\n",
|
| 303 |
+
" \"\"\"\n",
|
| 304 |
+
" rows = {\"prompt\": [], \"task\": [], \"seed\": []}\n",
|
| 305 |
+
" # TODO(conv-point):\n",
|
| 306 |
+
" # b3 = B3CortexFixedRouter(env=make_env())\n",
|
| 307 |
+
" # for ep in range(num_episodes):\n",
|
| 308 |
+
" # trajectory = b3.run_episode(task=\"outbreak_easy\", seed=ep)\n",
|
| 309 |
+
" # for router_step in trajectory.router_steps:\n",
|
| 310 |
+
" # rows[\"prompt\"].append(make_router_chat_prompt(router_step.metacognition_state))\n",
|
| 311 |
+
" # rows[\"task\"].append(\"outbreak_easy\")\n",
|
| 312 |
+
" # rows[\"seed\"].append(ep)\n",
|
| 313 |
+
" return (\n",
|
| 314 |
+
" Dataset.from_dict(rows)\n",
|
| 315 |
+
" if rows[\"prompt\"]\n",
|
| 316 |
+
" else Dataset.from_dict(\n",
|
| 317 |
+
" {\"prompt\": [\"placeholder until conv-point\"], \"task\": [\"outbreak_easy\"], \"seed\": [0]}\n",
|
| 318 |
+
" )\n",
|
| 319 |
+
" )\n",
|
| 320 |
+
"\n",
|
| 321 |
+
"\n",
|
| 322 |
+
"train_dataset = collect_b3_router_prompts()\n",
|
| 323 |
+
"print(f\"Dataset: {len(train_dataset)} rows (will be ~{TRAIN_EPISODES * 8} at conv-point)\")"
|
| 324 |
+
]
|
| 325 |
+
},
|
| 326 |
+
{
|
| 327 |
+
"cell_type": "markdown",
|
| 328 |
+
"id": "c3933fab20d04ec698c2621248eb3be0",
|
| 329 |
+
"metadata": {},
|
| 330 |
+
"source": [
|
| 331 |
+
"## 8. Reward function: full-episode rollout per (prompt, completion)\n",
|
| 332 |
+
"\n",
|
| 333 |
+
"Uses Phase-1 `outer_reward` summed across the episode, with Phase-2 `shape_reward` token-budget penalty.\n",
|
| 334 |
+
"\n",
|
| 335 |
+
"**TODO(conv-point):** the rollout loop calls `Council.step` with the trainable router policy. Until Session 13 lands, the loop is stubbed and returns 0.0 for every (prompt, completion)."
|
| 336 |
+
]
|
| 337 |
+
},
|
| 338 |
+
{
|
| 339 |
+
"cell_type": "code",
|
| 340 |
+
"execution_count": null,
|
| 341 |
+
"id": "4dd4641cc4064e0191573fe9c69df29b",
|
| 342 |
+
"metadata": {},
|
| 343 |
+
"outputs": [],
|
| 344 |
+
"source": [
|
| 345 |
+
"def cortex_router_reward(prompts, completions, task, seed, **_kwargs):\n",
|
| 346 |
+
" \"\"\"Reward = sum(per-tick obs.reward) over a full episode driven by the\n",
|
| 347 |
+
" trainable router's emitted RoutingAction JSON.\n",
|
| 348 |
+
"\n",
|
| 349 |
+
" Each (prompt, completion) pair represents ONE router decision; the\n",
|
| 350 |
+
" full episode reward is shared across all router decisions in that\n",
|
| 351 |
+
" episode (GRPO group-relative advantage handles credit assignment).\n",
|
| 352 |
+
"\n",
|
| 353 |
+
" TODO(conv-point): replace the stub body with the real Council-driven\n",
|
| 354 |
+
" rollout once Session 13 ships.\n",
|
| 355 |
+
" \"\"\"\n",
|
| 356 |
+
" rewards = []\n",
|
| 357 |
+
" for completion, t, s in zip(completions, task, seed):\n",
|
| 358 |
+
" # TODO(conv-point):\n",
|
| 359 |
+
" # try:\n",
|
| 360 |
+
" # routing_action = RoutingAction.model_validate_json(completion)\n",
|
| 361 |
+
" # except ValidationError:\n",
|
| 362 |
+
" # rewards.append(-1.0) # invalid JSON → terminal-equivalent penalty\n",
|
| 363 |
+
" # continue\n",
|
| 364 |
+
" # council = Council(routing_policy=trainable_router_from(routing_action),\n",
|
| 365 |
+
" # env=make_env(), brains=cortex_brains)\n",
|
| 366 |
+
" # episode_return = council.run_episode(task=t, seed=int(s)).total_reward\n",
|
| 367 |
+
" # rewards.append(episode_return)\n",
|
| 368 |
+
" rewards.append(0.0) # stub\n",
|
| 369 |
+
" return rewards"
|
| 370 |
+
]
|
| 371 |
+
},
|
| 372 |
+
{
|
| 373 |
+
"cell_type": "markdown",
|
| 374 |
+
"id": "8309879909854d7188b41380fd92a7c3",
|
| 375 |
+
"metadata": {},
|
| 376 |
+
"source": [
|
| 377 |
+
"## 9. GRPO training\n",
|
| 378 |
+
"\n",
|
| 379 |
+
"Same TRL `GRPOTrainer` shape as the B1 notebook. ~300 steps × group size 4 = ~1200 router decisions × full-episode rollouts. Wall-clock ~1.5 hours on a100-large at convergence point (dominated by brain-LLM rollouts, not router fine-tuning)."
|
| 380 |
+
]
|
| 381 |
+
},
|
| 382 |
+
{
|
| 383 |
+
"cell_type": "code",
|
| 384 |
+
"execution_count": null,
|
| 385 |
+
"id": "3ed186c9a28b402fb0bc4494df01f08d",
|
| 386 |
+
"metadata": {},
|
| 387 |
+
"outputs": [],
|
| 388 |
+
"source": [
|
| 389 |
+
"from trl import GRPOConfig, GRPOTrainer\n",
|
| 390 |
+
"\n",
|
| 391 |
+
"MAX_TRAIN_STEPS = 300\n",
|
| 392 |
+
"GROUP_SIZE = 4\n",
|
| 393 |
+
"MAX_PROMPT_LEN = 512\n",
|
| 394 |
+
"MAX_COMPLETION_LEN = 256 # router output is structured JSON (M-FR-11)\n",
|
| 395 |
+
"\n",
|
| 396 |
+
"training_args = GRPOConfig(\n",
|
| 397 |
+
" output_dir=\"/content/cortex_router_grpo_output\",\n",
|
| 398 |
+
" learning_rate=5e-6,\n",
|
| 399 |
+
" per_device_train_batch_size=GROUP_SIZE,\n",
|
| 400 |
+
" gradient_accumulation_steps=1,\n",
|
| 401 |
+
" num_generations=GROUP_SIZE,\n",
|
| 402 |
+
" max_prompt_length=MAX_PROMPT_LEN,\n",
|
| 403 |
+
" max_completion_length=MAX_COMPLETION_LEN,\n",
|
| 404 |
+
" max_steps=MAX_TRAIN_STEPS,\n",
|
| 405 |
+
" save_steps=100,\n",
|
| 406 |
+
" logging_steps=5,\n",
|
| 407 |
+
" report_to=\"none\",\n",
|
| 408 |
+
" bf16=True,\n",
|
| 409 |
+
" optim=\"adamw_8bit\",\n",
|
| 410 |
+
" temperature=0.8,\n",
|
| 411 |
+
" use_vllm=True,\n",
|
| 412 |
+
" vllm_mode=\"colocate\",\n",
|
| 413 |
+
" seed=42,\n",
|
| 414 |
+
")\n",
|
| 415 |
+
"\n",
|
| 416 |
+
"trainer = GRPOTrainer(\n",
|
| 417 |
+
" model=router_model,\n",
|
| 418 |
+
" processing_class=router_tokenizer,\n",
|
| 419 |
+
" reward_funcs=[cortex_router_reward],\n",
|
| 420 |
+
" args=training_args,\n",
|
| 421 |
+
" train_dataset=train_dataset,\n",
|
| 422 |
+
")\n",
|
| 423 |
+
"print(\"Router GRPOTrainer constructed.\")\n",
|
| 424 |
+
"print(\"# TODO(conv-point): uncomment trainer.train() once Cortex Session 13 lands.\")\n",
|
| 425 |
+
"# trainer.train()"
|
| 426 |
+
]
|
| 427 |
+
},
|
| 428 |
+
{
|
| 429 |
+
"cell_type": "markdown",
|
| 430 |
+
"id": "cb1e1581032b452c9409d6c6813c49d1",
|
| 431 |
+
"metadata": {},
|
| 432 |
+
"source": [
|
| 433 |
+
"## 10. Save the trained router LoRA to HF Hub"
|
| 434 |
+
]
|
| 435 |
+
},
|
| 436 |
+
{
|
| 437 |
+
"cell_type": "code",
|
| 438 |
+
"execution_count": null,
|
| 439 |
+
"id": "379cbbc1e968416e875cc15c1202d7eb",
|
| 440 |
+
"metadata": {},
|
| 441 |
+
"outputs": [],
|
| 442 |
+
"source": [
|
| 443 |
+
"from huggingface_hub import HfApi\n",
|
| 444 |
+
"\n",
|
| 445 |
+
"HUB_REPO = \"Angshuman28/crisisworld-cortex-router-llm\"\n",
|
| 446 |
+
"\n",
|
| 447 |
+
"router_model.save_pretrained(\"/content/cortex_router_lora\")\n",
|
| 448 |
+
"router_tokenizer.save_pretrained(\"/content/cortex_router_lora\")\n",
|
| 449 |
+
"\n",
|
| 450 |
+
"api = HfApi()\n",
|
| 451 |
+
"api.create_repo(HUB_REPO, exist_ok=True, repo_type=\"model\", private=False, token=HF_TOKEN)\n",
|
| 452 |
+
"api.upload_folder(\n",
|
| 453 |
+
" folder_path=\"/content/cortex_router_lora\",\n",
|
| 454 |
+
" repo_id=HUB_REPO,\n",
|
| 455 |
+
" repo_type=\"model\",\n",
|
| 456 |
+
" token=HF_TOKEN,\n",
|
| 457 |
+
")\n",
|
| 458 |
+
"print(f\"Saved to https://huggingface.co/{HUB_REPO}\")"
|
| 459 |
+
]
|
| 460 |
+
},
|
| 461 |
+
{
|
| 462 |
+
"cell_type": "markdown",
|
| 463 |
+
"id": "277c27b1587741f2af2001be3712ef0d",
|
| 464 |
+
"metadata": {},
|
| 465 |
+
"source": [
|
| 466 |
+
"## 11. Eval: B6 (trained LLM router) vs B3 (deterministic) on 3 tasks\n",
|
| 467 |
+
"\n",
|
| 468 |
+
"**TODO(conv-point):** the comparison loop below requires both B3CortexFixedRouter (deterministic) and a way to swap the trainable router into Council.routing_policy. The eval is the headline result for the convergence point — its sign decides whether B6 ships or B3 ships per the Phase 7 hard-exit gate.\n",
|
| 469 |
+
"\n",
|
| 470 |
+
"Decision rule (per spec): if reward over training steps is INCREASING, ship B6. Else ship B3."
|
| 471 |
+
]
|
| 472 |
+
},
|
| 473 |
+
{
|
| 474 |
+
"cell_type": "code",
|
| 475 |
+
"execution_count": null,
|
| 476 |
+
"id": "db7b79bc585a40fcaf58bf750017e135",
|
| 477 |
+
"metadata": {},
|
| 478 |
+
"outputs": [],
|
| 479 |
+
"source": [
|
| 480 |
+
"# TODO(conv-point):\n",
|
| 481 |
+
"# from cortex.routing_policy import DeterministicRouter, TrainableRouter\n",
|
| 482 |
+
"# from cortex.council import Council\n",
|
| 483 |
+
"# from baselines.cortex_fixed_router import B3CortexFixedRouter\n",
|
| 484 |
+
"#\n",
|
| 485 |
+
"# def run_b6_episode(task, seed):\n",
|
| 486 |
+
"# trainable = TrainableRouter(\n",
|
| 487 |
+
"# model=router_model,\n",
|
| 488 |
+
"# tokenizer=router_tokenizer,\n",
|
| 489 |
+
"# system_prompt=ROUTER_SYSTEM_PROMPT,\n",
|
| 490 |
+
"# featurize=metacog_state_to_prompt,\n",
|
| 491 |
+
"# )\n",
|
| 492 |
+
"# council = Council(routing_policy=trainable, env=make_env(), brains=cortex_brains)\n",
|
| 493 |
+
"# return council.run_episode(task=task, seed=seed).total_reward\n",
|
| 494 |
+
"#\n",
|
| 495 |
+
"# def run_b3_episode(task, seed):\n",
|
| 496 |
+
"# b3 = B3CortexFixedRouter(env=make_env())\n",
|
| 497 |
+
"# return b3.run_episode(task=task, seed=seed).total_reward\n",
|
| 498 |
+
"#\n",
|
| 499 |
+
"# b6_results = {t: run_b6_episode(t, 0) for t in TASKS}\n",
|
| 500 |
+
"# b3_results = {t: run_b3_episode(t, 0) for t in TASKS}\n",
|
| 501 |
+
"# print(f\"B6 (trained LLM router): {b6_results}\")\n",
|
| 502 |
+
"# print(f\"B3 (deterministic): {b3_results}\")\n",
|
| 503 |
+
"\n",
|
| 504 |
+
"print(\"Eval cell skeleton — uncomment when Cortex Session 13 ships\")"
|
| 505 |
+
]
|
| 506 |
+
},
|
| 507 |
+
{
|
| 508 |
+
"cell_type": "markdown",
|
| 509 |
+
"id": "916684f9a58a4a2aa5f864670399430d",
|
| 510 |
+
"metadata": {},
|
| 511 |
+
"source": [
|
| 512 |
+
"## 12. Plot training reward curve\n",
|
| 513 |
+
"\n",
|
| 514 |
+
"Phase 7 hard-exit gate watches this curve. If reward is increasing across training steps → ship B6. Else → ship B3."
|
| 515 |
+
]
|
| 516 |
+
},
|
| 517 |
+
{
|
| 518 |
+
"cell_type": "code",
|
| 519 |
+
"execution_count": null,
|
| 520 |
+
"id": "1671c31a24314836a5b85d7ef7fbf015",
|
| 521 |
+
"metadata": {},
|
| 522 |
+
"outputs": [],
|
| 523 |
+
"source": [
|
| 524 |
+
"import matplotlib.pyplot as plt\n",
|
| 525 |
+
"\n",
|
| 526 |
+
"log_path = \"/content/cortex_router_grpo_output/trainer_state.json\"\n",
|
| 527 |
+
"if os.path.exists(log_path):\n",
|
| 528 |
+
" import json as _json\n",
|
| 529 |
+
"\n",
|
| 530 |
+
" with open(log_path) as fh:\n",
|
| 531 |
+
" state = _json.load(fh)\n",
|
| 532 |
+
" history = state.get(\"log_history\", [])\n",
|
| 533 |
+
" rewards = [entry[\"reward\"] for entry in history if \"reward\" in entry]\n",
|
| 534 |
+
" if rewards:\n",
|
| 535 |
+
" fig, ax = plt.subplots(figsize=(8, 4))\n",
|
| 536 |
+
" ax.plot(rewards)\n",
|
| 537 |
+
" ax.set_xlabel(\"GRPO step\")\n",
|
| 538 |
+
" ax.set_ylabel(\"Mean reward\")\n",
|
| 539 |
+
" ax.set_title(\"Cortex small-LLM router — training reward\")\n",
|
| 540 |
+
" plt.tight_layout()\n",
|
| 541 |
+
" plt.show()\n",
|
| 542 |
+
" else:\n",
|
| 543 |
+
" print(\"No reward entries yet — uncomment trainer.train() at conv-point\")\n",
|
| 544 |
+
"else:\n",
|
| 545 |
+
" print(\"No trainer state yet — uncomment trainer.train() at conv-point\")"
|
| 546 |
+
]
|
| 547 |
+
}
|
| 548 |
+
],
|
| 549 |
+
"metadata": {
|
| 550 |
+
"accelerator": "GPU",
|
| 551 |
+
"colab": {
|
| 552 |
+
"gpuType": "A100",
|
| 553 |
+
"provenance": []
|
| 554 |
+
},
|
| 555 |
+
"kernelspec": {
|
| 556 |
+
"display_name": "Python 3",
|
| 557 |
+
"language": "python",
|
| 558 |
+
"name": "python3"
|
| 559 |
+
},
|
| 560 |
+
"language_info": {
|
| 561 |
+
"name": "python",
|
| 562 |
+
"version": "3.10"
|
| 563 |
+
}
|
| 564 |
+
},
|
| 565 |
+
"nbformat": 4,
|
| 566 |
+
"nbformat_minor": 5
|
| 567 |
+
}
|
openenv.yaml
CHANGED
|
@@ -1,7 +1,7 @@
|
|
| 1 |
-
spec_version: 1
|
| 2 |
-
name: CrisisWorldCortex
|
| 3 |
-
type: space
|
| 4 |
-
runtime: fastapi
|
| 5 |
-
app: server.app:app
|
| 6 |
-
port: 8000
|
| 7 |
-
|
|
|
|
| 1 |
+
spec_version: 1
|
| 2 |
+
name: CrisisWorldCortex
|
| 3 |
+
type: space
|
| 4 |
+
runtime: fastapi
|
| 5 |
+
app: server.app:app
|
| 6 |
+
port: 8000
|
| 7 |
+
|
openenv_CrisisWorldCortex.egg-info/PKG-INFO
ADDED
|
@@ -0,0 +1,13 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
Metadata-Version: 2.4
|
| 2 |
+
Name: openenv-CrisisWorldCortex
|
| 3 |
+
Version: 0.1.0
|
| 4 |
+
Summary: Crisisworldcortex environment for OpenEnv
|
| 5 |
+
Requires-Python: >=3.10
|
| 6 |
+
Requires-Dist: openenv-core[core]==0.2.3
|
| 7 |
+
Requires-Dist: openai<3.0,>=2.0
|
| 8 |
+
Requires-Dist: python-dotenv>=1.0.0
|
| 9 |
+
Provides-Extra: dev
|
| 10 |
+
Requires-Dist: pre-commit>=4.0.0; extra == "dev"
|
| 11 |
+
Requires-Dist: pytest>=8.0.0; extra == "dev"
|
| 12 |
+
Requires-Dist: pytest-cov>=4.0.0; extra == "dev"
|
| 13 |
+
Requires-Dist: ruff>=0.8.0; extra == "dev"
|
openenv_CrisisWorldCortex.egg-info/SOURCES.txt
ADDED
|
@@ -0,0 +1,48 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
README.md
|
| 2 |
+
__init__.py
|
| 3 |
+
client.py
|
| 4 |
+
inference.py
|
| 5 |
+
models.py
|
| 6 |
+
pyproject.toml
|
| 7 |
+
./__init__.py
|
| 8 |
+
./client.py
|
| 9 |
+
./inference.py
|
| 10 |
+
./models.py
|
| 11 |
+
openenv_CrisisWorldCortex.egg-info/PKG-INFO
|
| 12 |
+
openenv_CrisisWorldCortex.egg-info/SOURCES.txt
|
| 13 |
+
openenv_CrisisWorldCortex.egg-info/dependency_links.txt
|
| 14 |
+
openenv_CrisisWorldCortex.egg-info/entry_points.txt
|
| 15 |
+
openenv_CrisisWorldCortex.egg-info/requires.txt
|
| 16 |
+
openenv_CrisisWorldCortex.egg-info/top_level.txt
|
| 17 |
+
server/CrisisWorldCortex_environment.py
|
| 18 |
+
server/__init__.py
|
| 19 |
+
server/app.py
|
| 20 |
+
tests/test_actions_round_trip.py
|
| 21 |
+
tests/test_baseline_b1.py
|
| 22 |
+
tests/test_baseline_b2.py
|
| 23 |
+
tests/test_cortex_brain_executive.py
|
| 24 |
+
tests/test_cortex_brain_smoke.py
|
| 25 |
+
tests/test_cortex_lenses.py
|
| 26 |
+
tests/test_cortex_perception.py
|
| 27 |
+
tests/test_cortex_subagents.py
|
| 28 |
+
tests/test_env_reset_kwargs.py
|
| 29 |
+
tests/test_env_step_reward_wiring.py
|
| 30 |
+
tests/test_import_graph.py
|
| 31 |
+
tests/test_legal_constraint_enforcement.py
|
| 32 |
+
tests/test_llm_client.py
|
| 33 |
+
tests/test_observation_no_latent_leak.py
|
| 34 |
+
tests/test_outer_reward_in_range.py
|
| 35 |
+
tests/test_outer_reward_non_constancy.py
|
| 36 |
+
tests/test_outer_reward_terminal_bonus.py
|
| 37 |
+
tests/test_package_exports.py
|
| 38 |
+
tests/test_reward_signal_quality.py
|
| 39 |
+
tests/test_schemas_roundtrip.py
|
| 40 |
+
tests/test_simulator_determinism.py
|
| 41 |
+
tests/test_simulator_random_episode.py
|
| 42 |
+
tests/test_simulator_task_configs.py
|
| 43 |
+
tests/test_smoke_env.py
|
| 44 |
+
tests/test_stdout_format.py
|
| 45 |
+
tests/test_synthetic_rejection_payload.py
|
| 46 |
+
tests/test_training_eval_metrics.py
|
| 47 |
+
tests/test_training_reward_shaping.py
|
| 48 |
+
tests/test_training_rollout_buffer.py
|
openenv_CrisisWorldCortex.egg-info/dependency_links.txt
ADDED
|
@@ -0,0 +1 @@
|
|
|
|
|
|
|
| 1 |
+
|
openenv_CrisisWorldCortex.egg-info/entry_points.txt
ADDED
|
@@ -0,0 +1,2 @@
|
|
|
|
|
|
|
|
|
|
| 1 |
+
[console_scripts]
|
| 2 |
+
server = CrisisWorldCortex.server.app:main
|
openenv_CrisisWorldCortex.egg-info/requires.txt
ADDED
|
@@ -0,0 +1,9 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
openenv-core[core]==0.2.3
|
| 2 |
+
openai<3.0,>=2.0
|
| 3 |
+
python-dotenv>=1.0.0
|
| 4 |
+
|
| 5 |
+
[dev]
|
| 6 |
+
pre-commit>=4.0.0
|
| 7 |
+
pytest>=8.0.0
|
| 8 |
+
pytest-cov>=4.0.0
|
| 9 |
+
ruff>=0.8.0
|
openenv_CrisisWorldCortex.egg-info/top_level.txt
ADDED
|
@@ -0,0 +1 @@
|
|
|
|
|
|
|
| 1 |
+
CrisisWorldCortex
|
pyproject.toml
CHANGED
|
@@ -25,6 +25,7 @@ dependencies = [
|
|
| 25 |
# future uv sync from silently pulling 3.x. Bump explicitly when 3.0
|
| 26 |
# ships and we've verified compatibility.
|
| 27 |
"openai>=2.0,<3.0",
|
|
|
|
| 28 |
]
|
| 29 |
|
| 30 |
[project.optional-dependencies]
|
|
|
|
| 25 |
# future uv sync from silently pulling 3.x. Bump explicitly when 3.0
|
| 26 |
# ships and we've verified compatibility.
|
| 27 |
"openai>=2.0,<3.0",
|
| 28 |
+
"python-dotenv>=1.0.0"
|
| 29 |
]
|
| 30 |
|
| 31 |
[project.optional-dependencies]
|
server/CLAUDE.md
CHANGED
|
@@ -15,33 +15,24 @@ Cortex code (brains, subagents, router, metacognition), LLM clients, training lo
|
|
| 15 |
|
| 16 |
## Import rule (binding)
|
| 17 |
|
| 18 |
-
**Inside `server/`, use
|
| 19 |
-
`from .simulator import ...` from one-level modules and
|
| 20 |
-
`from ..simulator import ...` from nested modules. This works when the
|
| 21 |
-
server is loaded as either `CrisisWorldCortex.server.*` (`uv run server`)
|
| 22 |
-
or top-level `server.*` (Docker / `uvicorn server.app:app`). Never use
|
| 23 |
-
`from CrisisWorldCortex.server...` for server-internal imports.
|
| 24 |
|
| 25 |
-
For `models`, use the
|
| 26 |
|
| 27 |
```python
|
| 28 |
-
|
|
|
|
|
|
|
|
|
|
| 29 |
```
|
| 30 |
|
| 31 |
-
The server runs under
|
| 32 |
-
`uvicorn server.app:app`, Docker `cd /app/env && uvicorn server.app:app`).
|
| 33 |
-
Canonical wire imports keep Pydantic model identity stable across those modes.
|
| 34 |
|
| 35 |
-
**Wire-type imports from deep modules (binding)**:
|
| 36 |
-
files (`server/simulator/*`, `server/graders/*`) must still use
|
| 37 |
-
`from CrisisWorldCortex.models import …`; `..models` from those depths
|
| 38 |
-
resolves to a non-existent `CrisisWorldCortex.server.models`, and a bare
|
| 39 |
-
fallback loads a second `models` module.
|
| 40 |
|
| 41 |
## Allowed imports
|
| 42 |
|
| 43 |
-
`
|
| 44 |
-
`server/graders/*`, `openenv.core.*`, stdlib, FastAPI, Pydantic, numpy.
|
| 45 |
|
| 46 |
## Forbidden imports
|
| 47 |
|
|
|
|
| 15 |
|
| 16 |
## Import rule (binding)
|
| 17 |
|
| 18 |
+
**Inside `server/`, use `from server.simulator import …`, never `from CrisisWorldCortex.server.simulator import …`.** Same rule for `server.graders`.
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 19 |
|
| 20 |
+
For `models`, use the dual-import fallback in every new server module that needs it:
|
| 21 |
|
| 22 |
```python
|
| 23 |
+
try:
|
| 24 |
+
from ..models import CrisisworldcortexAction, CrisisworldcortexObservation
|
| 25 |
+
except (ImportError, ModuleNotFoundError):
|
| 26 |
+
from models import CrisisworldcortexAction, CrisisworldcortexObservation
|
| 27 |
```
|
| 28 |
|
| 29 |
+
The server runs under ≥ 3 import contexts (`python -m server.app`, `uvicorn server.app:app`, Docker `cd /app/env && uvicorn server.app:app`); the fallback is load-bearing.
|
|
|
|
|
|
|
| 30 |
|
| 31 |
+
**Wire-type imports from deep modules (binding)**: the dual-import fallback above only works for files **one level** inside `server/` (`server/CrisisWorldCortex_environment.py`, `server/app.py`). Two-or-more-levels-deep files (`server/simulator/*`, `server/graders/*`) **must** use `from CrisisWorldCortex.models import …` directly — `..models` from those depths resolves to a non-existent `CrisisWorldCortex.server.models`, the fallback fires, and bare `models` loads as a separate `sys.modules` entry, breaking Pydantic discriminator validation against types imported via the canonical path. Session 5a's `server/simulator/seir_model.py` and `server/simulator/tasks.py` document this with inline comments.
|
|
|
|
|
|
|
|
|
|
|
|
|
| 32 |
|
| 33 |
## Allowed imports
|
| 34 |
|
| 35 |
+
`models` (via dual-import fallback), `openenv.core.*`, `server/simulator/*`, `server/graders/*`, stdlib, FastAPI, Pydantic, numpy.
|
|
|
|
| 36 |
|
| 37 |
## Forbidden imports
|
| 38 |
|
server/CrisisWorldCortex_environment.py
CHANGED
|
@@ -121,10 +121,26 @@ class CrisisworldcortexEnvironment(Environment):
|
|
| 121 |
)
|
| 122 |
self._state.step_count += 1
|
| 123 |
self._world_state = apply_tick(self._world_state, action.action)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 124 |
obs = make_observation(self._world_state)
|
| 125 |
-
# Per design §15: r_outer is the only env-side reward signal, in
|
| 126 |
-
# Terminal bonus (+/-0.20)
|
| 127 |
-
# design §14.3 — never
|
|
|
|
| 128 |
obs.reward = outer_reward(self._world_state, action.action)
|
| 129 |
return obs
|
| 130 |
|
|
|
|
| 121 |
)
|
| 122 |
self._state.step_count += 1
|
| 123 |
self._world_state = apply_tick(self._world_state, action.action)
|
| 124 |
+
# Parse-failure terminal contract (design §19, Phase-1 restoration):
|
| 125 |
+
# the synthetic parse_failure_marker (PublicCommunication with
|
| 126 |
+
# honesty=0.0, magic-string discriminator per Phase-A M3-B) ends
|
| 127 |
+
# the episode as state.terminal = "failure". apply_tick may have
|
| 128 |
+
# set terminal to "none"/"success"/"timeout" via the SEIR rules;
|
| 129 |
+
# we override here because parse-failure is a harness-level fault,
|
| 130 |
+
# not a simulator-level event.
|
| 131 |
+
payload = action.action
|
| 132 |
+
if (
|
| 133 |
+
payload.kind == "public_communication"
|
| 134 |
+
and getattr(payload, "honesty", None) == 0.0
|
| 135 |
+
and self._world_state.recent_action_log
|
| 136 |
+
and not self._world_state.recent_action_log[-1].accepted
|
| 137 |
+
):
|
| 138 |
+
self._world_state.terminal = "failure"
|
| 139 |
obs = make_observation(self._world_state)
|
| 140 |
+
# Per design §15: r_outer is the only env-side reward signal, in
|
| 141 |
+
# [-1.0, 1.0] post-Phase-1 (was [0, 1]). Terminal bonus (+/-0.20)
|
| 142 |
+
# is composed downstream by the trainer per design §14.3 — never
|
| 143 |
+
# bundled into obs.reward.
|
| 144 |
obs.reward = outer_reward(self._world_state, action.action)
|
| 145 |
return obs
|
| 146 |
|
server/__init__.py
CHANGED
|
@@ -1,11 +1,11 @@
|
|
| 1 |
-
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
| 2 |
-
# All rights reserved.
|
| 3 |
-
#
|
| 4 |
-
# This source code is licensed under the BSD-style license found in the
|
| 5 |
-
# LICENSE file in the root directory of this source tree.
|
| 6 |
-
|
| 7 |
-
"""Crisisworldcortex environment server components."""
|
| 8 |
-
|
| 9 |
-
from .CrisisWorldCortex_environment import CrisisworldcortexEnvironment
|
| 10 |
-
|
| 11 |
-
__all__ = ["CrisisworldcortexEnvironment"]
|
|
|
|
| 1 |
+
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
| 2 |
+
# All rights reserved.
|
| 3 |
+
#
|
| 4 |
+
# This source code is licensed under the BSD-style license found in the
|
| 5 |
+
# LICENSE file in the root directory of this source tree.
|
| 6 |
+
|
| 7 |
+
"""Crisisworldcortex environment server components."""
|
| 8 |
+
|
| 9 |
+
from .CrisisWorldCortex_environment import CrisisworldcortexEnvironment
|
| 10 |
+
|
| 11 |
+
__all__ = ["CrisisworldcortexEnvironment"]
|
server/app.py
CHANGED
|
@@ -29,8 +29,8 @@ Usage:
|
|
| 29 |
"""
|
| 30 |
|
| 31 |
try:
|
| 32 |
-
|
| 33 |
-
from openenv.core.env_server import create_web_interface_app as create_app
|
| 34 |
except Exception as e: # pragma: no cover
|
| 35 |
raise ImportError(
|
| 36 |
"openenv is required for the web interface. Install dependencies with '\n uv sync\n'"
|
|
|
|
| 29 |
"""
|
| 30 |
|
| 31 |
try:
|
| 32 |
+
from openenv.core.env_server.http_server import create_app
|
| 33 |
+
# from openenv.core.env_server import create_web_interface_app as create_app
|
| 34 |
except Exception as e: # pragma: no cover
|
| 35 |
raise ImportError(
|
| 36 |
"openenv is required for the web interface. Install dependencies with '\n uv sync\n'"
|
server/graders/outer_reward.py
CHANGED
|
@@ -5,16 +5,29 @@
|
|
| 5 |
# LICENSE file in the root directory of this source tree.
|
| 6 |
|
| 7 |
"""
|
| 8 |
-
Outer reward grader for CrisisWorld (design §15).
|
| 9 |
|
| 10 |
Public API (re-exported via ``server/graders/__init__.py``):
|
| 11 |
- ``outer_reward(state, action) -> float`` — 6-component weighted score in
|
| 12 |
-
``[
|
| 13 |
-
|
|
|
|
|
|
|
|
|
|
| 14 |
- ``terminal_bonus(state) -> float`` — episode-end ±0.20 / 0.0 bonus,
|
| 15 |
composed by trainer in ``training/reward_shaping.py`` per design §14.3
|
| 16 |
-
(``episode_return = Σ_t r_outer + terminal_bonus``). Kept separate
|
| 17 |
-
per-tick
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 18 |
|
| 19 |
Wire-protocol imports use the absolute path ``CrisisWorldCortex.models``
|
| 20 |
because this file lives two levels deep inside ``server/`` — see
|
|
@@ -38,25 +51,42 @@ from ..simulator import (
|
|
| 38 |
)
|
| 39 |
|
| 40 |
# ============================================================================
|
| 41 |
-
# Component weights (design §15; sum to 1.00)
|
| 42 |
# ============================================================================
|
| 43 |
|
| 44 |
-
W_INFECT = 0.35
|
| 45 |
-
W_TIME = 0.18
|
| 46 |
-
W_HOSP = 0.17
|
| 47 |
-
W_CASC = 0.15
|
| 48 |
-
W_POLICY = 0.12
|
| 49 |
-
W_FAIR = 0.03
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 50 |
|
| 51 |
# Terminal-bonus magnitudes (design §14.3 / §15).
|
| 52 |
TERMINAL_BONUS_SUCCESS = 0.20
|
| 53 |
TERMINAL_BONUS_FAILURE = -0.20
|
| 54 |
|
| 55 |
-
# r_casc threshold:
|
| 56 |
-
#
|
| 57 |
-
#
|
| 58 |
R_CASC_HOT_THRESHOLD = CATASTROPHIC_INFECTION_THRESHOLD
|
| 59 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 60 |
|
| 61 |
def _clamp01(x: float) -> float:
|
| 62 |
if x < 0.0:
|
|
@@ -76,34 +106,56 @@ def _hospital_load(region_I: float) -> float:
|
|
| 76 |
return _clamp01(region_I * HOSPITALIZATION_FRACTION_OF_I / HOSPITAL_CAPACITY_FRACTION)
|
| 77 |
|
| 78 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 79 |
def outer_reward(state: WorldState, action: OuterActionPayload) -> float:
|
| 80 |
-
"""Compute per-tick outer reward in ``[
|
| 81 |
|
| 82 |
Read post-``apply_tick`` state: ``state.regions[*].I`` is the
|
| 83 |
just-stepped ground-truth infection fraction. ``recent_action_log[-1]``
|
| 84 |
holds the just-dispatched action's acceptance flag.
|
| 85 |
|
| 86 |
-
Six components (design §15):
|
| 87 |
-
r_infect = 1 - mean(I)
|
| 88 |
-
r_time = 1 - tick / max_ticks
|
| 89 |
-
r_hosp = 1 - mean(hospital_load)
|
| 90 |
-
r_casc = 1
|
| 91 |
-
r_policy = 1
|
| 92 |
-
r_fair = 1 - pstdev(I)
|
| 93 |
|
| 94 |
The ``action`` argument is the action just dispatched. We read its
|
| 95 |
acceptance flag from ``state.recent_action_log[-1]`` rather than
|
| 96 |
re-dispatching — the simulator already recorded it, and re-dispatch
|
| 97 |
-
would mutate state.
|
| 98 |
-
parity with ``training_reward(trajectory)`` and to support future
|
| 99 |
-
action-shape-dependent shaping without a signature break.
|
| 100 |
"""
|
| 101 |
if not state.regions:
|
| 102 |
return 0.0
|
| 103 |
|
| 104 |
I_values = [r.I for r in state.regions]
|
|
|
|
| 105 |
|
| 106 |
-
r_infect
|
|
|
|
| 107 |
|
| 108 |
if state.max_ticks > 0:
|
| 109 |
r_time = _clamp01(1.0 - state.tick / state.max_ticks)
|
|
@@ -111,16 +163,22 @@ def outer_reward(state: WorldState, action: OuterActionPayload) -> float:
|
|
| 111 |
r_time = 0.0
|
| 112 |
|
| 113 |
hosp_loads = [_hospital_load(I) for I in I_values]
|
| 114 |
-
|
|
|
|
| 115 |
|
| 116 |
-
|
| 117 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
| 118 |
|
|
|
|
| 119 |
if state.recent_action_log:
|
| 120 |
last_entry = state.recent_action_log[-1]
|
| 121 |
-
r_policy =
|
| 122 |
else:
|
| 123 |
-
r_policy =
|
| 124 |
|
| 125 |
if len(I_values) >= 2:
|
| 126 |
r_fair = _clamp01(1.0 - statistics.pstdev(I_values))
|
|
@@ -135,8 +193,12 @@ def outer_reward(state: WorldState, action: OuterActionPayload) -> float:
|
|
| 135 |
+ W_POLICY * r_policy
|
| 136 |
+ W_FAIR * r_fair
|
| 137 |
)
|
| 138 |
-
# Final clamp
|
| 139 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
| 140 |
|
| 141 |
|
| 142 |
def terminal_bonus(state: WorldState) -> float:
|
|
|
|
| 5 |
# LICENSE file in the root directory of this source tree.
|
| 6 |
|
| 7 |
"""
|
| 8 |
+
Outer reward grader for CrisisWorld (design §15 + §19, post-Phase-1 fix).
|
| 9 |
|
| 10 |
Public API (re-exported via ``server/graders/__init__.py``):
|
| 11 |
- ``outer_reward(state, action) -> float`` — 6-component weighted score in
|
| 12 |
+
``[-1.0, 1.0]`` (Workstream-B Phase-1 range relaxation; was
|
| 13 |
+
``[0.0, 1.0]``). Negative values come from ``r_policy = -0.5`` on rejected
|
| 14 |
+
actions and ``r_policy = -1.0`` on parse-failure markers (per design §19).
|
| 15 |
+
Computed on post-``apply_tick`` state. Still the only env-side reward
|
| 16 |
+
signal per ``server/CLAUDE.md``.
|
| 17 |
- ``terminal_bonus(state) -> float`` — episode-end ±0.20 / 0.0 bonus,
|
| 18 |
composed by trainer in ``training/reward_shaping.py`` per design §14.3
|
| 19 |
+
(``episode_return = Σ_t r_outer + terminal_bonus``). Kept separate from
|
| 20 |
+
the per-tick scalar.
|
| 21 |
+
|
| 22 |
+
Phase-1 changes (Workstream B):
|
| 23 |
+
- Steeper sensitivity on ``r_infect`` (``× 20``) and ``r_hosp`` (``× 10``)
|
| 24 |
+
so the gentle outbreak_easy task still produces gradient.
|
| 25 |
+
- Continuous ``r_casc`` (``1 - max(I)/0.30`` clamped) replaces binary.
|
| 26 |
+
- ``r_policy`` ∈ {-1.0 (parse-failure), -0.5 (rejected), 0.0 (accepted
|
| 27 |
+
no_op), +1.0 (accepted real action)} restoring §19 magnitudes.
|
| 28 |
+
- Weight redistribution: W_POLICY 0.12 → 0.35 (signal-driver); W_TIME
|
| 29 |
+
0.18 → 0.05 (action-independent noise); other components rebalanced.
|
| 30 |
+
- Final ``[0,1]`` clamp dropped (Phase-A M2-A).
|
| 31 |
|
| 32 |
Wire-protocol imports use the absolute path ``CrisisWorldCortex.models``
|
| 33 |
because this file lives two levels deep inside ``server/`` — see
|
|
|
|
| 51 |
)
|
| 52 |
|
| 53 |
# ============================================================================
|
| 54 |
+
# Component weights (design §15 + Phase-1 redistribution; sum to 1.00)
|
| 55 |
# ============================================================================
|
| 56 |
|
| 57 |
+
W_INFECT = 0.15 # was 0.35; iter-1 reduced because outbreak_easy keeps mean(I) tiny
|
| 58 |
+
W_TIME = 0.05 # was 0.18; action-independent, no signal value
|
| 59 |
+
W_HOSP = 0.10 # was 0.17; iter-1 reduction (gentle env keeps hosp_load low)
|
| 60 |
+
W_CASC = 0.10 # was 0.15; iter-1 reduction
|
| 61 |
+
W_POLICY = 0.55 # was 0.12; iter-1 dominant signal driver — accepted-real vs no_op vs rejected
|
| 62 |
+
W_FAIR = 0.05 # was 0.03; tiny boost
|
| 63 |
+
|
| 64 |
+
# Steepness coefficients (Phase-A M6, ONE-iteration tentative).
|
| 65 |
+
# r_infect ≈ 0 when mean(I) >= 0.05; near 1 when mean(I) <= 0.0 → strong gradient.
|
| 66 |
+
# r_hosp similarly sensitive to mean hospital_load.
|
| 67 |
+
R_INFECT_STEEPNESS = 20.0
|
| 68 |
+
R_HOSP_STEEPNESS = 10.0
|
| 69 |
+
|
| 70 |
+
# r_policy values per design §19 (Phase-1 restoration).
|
| 71 |
+
R_POLICY_PARSE_FAILURE = -1.0 # synthetic parse-failure marker
|
| 72 |
+
R_POLICY_REJECTED = -0.5 # well-formed-illegal (V2 / legal-violation)
|
| 73 |
+
R_POLICY_NOOP_ACCEPTED = 0.0 # accepted no-op (legal but inactive)
|
| 74 |
+
R_POLICY_REAL_ACCEPTED = 1.0 # accepted real intervention
|
| 75 |
|
| 76 |
# Terminal-bonus magnitudes (design §14.3 / §15).
|
| 77 |
TERMINAL_BONUS_SUCCESS = 0.20
|
| 78 |
TERMINAL_BONUS_FAILURE = -0.20
|
| 79 |
|
| 80 |
+
# r_casc threshold: at max(I) >= this, r_casc = 0 (catastrophe imminent).
|
| 81 |
+
# Matches design §6.4's catastrophic-infection threshold so cascade
|
| 82 |
+
# signal aligns with the failure terminal.
|
| 83 |
R_CASC_HOT_THRESHOLD = CATASTROPHIC_INFECTION_THRESHOLD
|
| 84 |
|
| 85 |
+
# Magic-string discriminator for parse-failure marker (Phase-A M3-B):
|
| 86 |
+
# baselines.flat_agent.parse_failure_marker emits PublicCommunication with
|
| 87 |
+
# honesty=0.0; intentional V2 attempts use honesty > 0.0.
|
| 88 |
+
PARSE_FAILURE_HONESTY_SENTINEL = 0.0
|
| 89 |
+
|
| 90 |
|
| 91 |
def _clamp01(x: float) -> float:
|
| 92 |
if x < 0.0:
|
|
|
|
| 106 |
return _clamp01(region_I * HOSPITALIZATION_FRACTION_OF_I / HOSPITAL_CAPACITY_FRACTION)
|
| 107 |
|
| 108 |
|
| 109 |
+
def _r_policy_value(action: OuterActionPayload, accepted: bool) -> float:
|
| 110 |
+
"""Compute ``r_policy`` per design §19 four-state contract.
|
| 111 |
+
|
| 112 |
+
Returns one of {-1.0, -0.5, 0.0, +1.0} based on (action.kind, accepted).
|
| 113 |
+
Parse-failure detection uses the ``honesty == 0.0`` sentinel on a
|
| 114 |
+
rejected ``PublicCommunication`` payload (Phase-A M3-B magic string).
|
| 115 |
+
"""
|
| 116 |
+
if not accepted:
|
| 117 |
+
# Rejected branch. Distinguish parse-failure marker from intentional
|
| 118 |
+
# V2-PublicCommunication / legal-violation rejection.
|
| 119 |
+
if (
|
| 120 |
+
action.kind == "public_communication"
|
| 121 |
+
and getattr(action, "honesty", None) == PARSE_FAILURE_HONESTY_SENTINEL
|
| 122 |
+
):
|
| 123 |
+
return R_POLICY_PARSE_FAILURE
|
| 124 |
+
return R_POLICY_REJECTED
|
| 125 |
+
# Accepted branch.
|
| 126 |
+
if action.kind == "no_op":
|
| 127 |
+
return R_POLICY_NOOP_ACCEPTED
|
| 128 |
+
return R_POLICY_REAL_ACCEPTED
|
| 129 |
+
|
| 130 |
+
|
| 131 |
def outer_reward(state: WorldState, action: OuterActionPayload) -> float:
|
| 132 |
+
"""Compute per-tick outer reward in ``[-1.0, 1.0]`` (post-Phase-1 range).
|
| 133 |
|
| 134 |
Read post-``apply_tick`` state: ``state.regions[*].I`` is the
|
| 135 |
just-stepped ground-truth infection fraction. ``recent_action_log[-1]``
|
| 136 |
holds the just-dispatched action's acceptance flag.
|
| 137 |
|
| 138 |
+
Six components (design §15 + Phase-1 fix):
|
| 139 |
+
r_infect = max(0, 1 - 20 × mean(I)) # weight 0.25
|
| 140 |
+
r_time = 1 - tick / max_ticks # weight 0.05
|
| 141 |
+
r_hosp = max(0, 1 - 10 × mean(hospital_load)) # weight 0.15
|
| 142 |
+
r_casc = max(0, 1 - max(I) / 0.30) # weight 0.15
|
| 143 |
+
r_policy = {-1.0, -0.5, 0.0, +1.0} per §19 # weight 0.35
|
| 144 |
+
r_fair = 1 - pstdev(I) # weight 0.05
|
| 145 |
|
| 146 |
The ``action`` argument is the action just dispatched. We read its
|
| 147 |
acceptance flag from ``state.recent_action_log[-1]`` rather than
|
| 148 |
re-dispatching — the simulator already recorded it, and re-dispatch
|
| 149 |
+
would mutate state.
|
|
|
|
|
|
|
| 150 |
"""
|
| 151 |
if not state.regions:
|
| 152 |
return 0.0
|
| 153 |
|
| 154 |
I_values = [r.I for r in state.regions]
|
| 155 |
+
mean_I = sum(I_values) / len(I_values)
|
| 156 |
|
| 157 |
+
# r_infect: steepened so gentle outbreak_easy still produces gradient.
|
| 158 |
+
r_infect = _clamp01(1.0 - R_INFECT_STEEPNESS * mean_I)
|
| 159 |
|
| 160 |
if state.max_ticks > 0:
|
| 161 |
r_time = _clamp01(1.0 - state.tick / state.max_ticks)
|
|
|
|
| 163 |
r_time = 0.0
|
| 164 |
|
| 165 |
hosp_loads = [_hospital_load(I) for I in I_values]
|
| 166 |
+
mean_hosp = sum(hosp_loads) / len(hosp_loads)
|
| 167 |
+
r_hosp = _clamp01(1.0 - R_HOSP_STEEPNESS * mean_hosp)
|
| 168 |
|
| 169 |
+
# r_casc: continuous ramp (1.0 at max(I)=0; 0.0 at max(I) >= threshold).
|
| 170 |
+
max_I = max(I_values)
|
| 171 |
+
if R_CASC_HOT_THRESHOLD > 0:
|
| 172 |
+
r_casc = _clamp01(1.0 - max_I / R_CASC_HOT_THRESHOLD)
|
| 173 |
+
else:
|
| 174 |
+
r_casc = 0.0
|
| 175 |
|
| 176 |
+
# r_policy: design §19 four-state contract.
|
| 177 |
if state.recent_action_log:
|
| 178 |
last_entry = state.recent_action_log[-1]
|
| 179 |
+
r_policy = _r_policy_value(last_entry.action, last_entry.accepted)
|
| 180 |
else:
|
| 181 |
+
r_policy = R_POLICY_REAL_ACCEPTED # No action dispatched yet → no penalty.
|
| 182 |
|
| 183 |
if len(I_values) >= 2:
|
| 184 |
r_fair = _clamp01(1.0 - statistics.pstdev(I_values))
|
|
|
|
| 193 |
+ W_POLICY * r_policy
|
| 194 |
+ W_FAIR * r_fair
|
| 195 |
)
|
| 196 |
+
# Final clamp to [-1.0, 1.0] (no longer [0, 1] — M2-A drops the floor).
|
| 197 |
+
if score < -1.0:
|
| 198 |
+
return -1.0
|
| 199 |
+
if score > 1.0:
|
| 200 |
+
return 1.0
|
| 201 |
+
return score
|
| 202 |
|
| 203 |
|
| 204 |
def terminal_bonus(state: WorldState) -> float:
|
server/requirements.txt
CHANGED
|
@@ -1,6 +1,6 @@
|
|
| 1 |
-
openenv[core]>=0.2.0
|
| 2 |
-
fastapi>=0.115.0
|
| 3 |
-
uvicorn>=0.24.0
|
| 4 |
-
|
| 5 |
-
|
| 6 |
-
|
|
|
|
| 1 |
+
openenv[core]>=0.2.0
|
| 2 |
+
fastapi>=0.115.0
|
| 3 |
+
uvicorn>=0.24.0
|
| 4 |
+
|
| 5 |
+
|
| 6 |
+
|
tests/CLAUDE.md
CHANGED
|
@@ -1,48 +1,48 @@
|
|
| 1 |
-
# tests/CLAUDE.md
|
| 2 |
-
|
| 3 |
-
Test surface per subsystem. Smoke bar first, boundary tests next, coverage last.
|
| 4 |
-
|
| 5 |
-
## Belongs here
|
| 6 |
-
|
| 7 |
-
- `conftest.py` — repo root on `sys.path` for bare-name imports.
|
| 8 |
-
- One test module per subsystem boundary (table below).
|
| 9 |
-
|
| 10 |
-
## Does not belong here
|
| 11 |
-
|
| 12 |
-
Helpers that mutate real graders, simulator state, or disk. Fixtures that hit the live HF Space — mock or `pytest.skip`.
|
| 13 |
-
|
| 14 |
-
## Run commands
|
| 15 |
-
|
| 16 |
-
```bash
|
| 17 |
-
uv run python -m pytest tests/ -v # all
|
| 18 |
-
uv run python -m pytest tests/test_smoke_env.py::test_reset_returns_valid_observation -v # one
|
| 19 |
-
uv run python -m pytest --cov tests/ # coverage
|
| 20 |
-
```
|
| 21 |
-
|
| 22 |
-
## Required tests — each maps to exactly one subsystem contract
|
| 23 |
-
|
| 24 |
-
| File | Scope | Asserts |
|
| 25 |
-
|---|---|---|
|
| 26 |
-
| `test_package_exports.py` | wire package | Root `__init__` re-exports `CrisisworldcortexAction/Observation/Env`. |
|
| 27 |
-
| `test_smoke_env.py` | `server/` env | `reset()` / `step()` return a valid `CrisisworldcortexObservation`. |
|
| 28 |
-
| `test_actions_round_trip.py` | `server/` env | 6 MVP outer actions round-trip; `public_communication` is rejected at runtime. |
|
| 29 |
-
| `test_reward_shape.py` | `server/graders/` | Every grader returns values in `[0.0, 1.0]`. |
|
| 30 |
-
| `test_reward_non_constancy.py` | `server/graders/` | Grader output varies across ≥ 2 synthetic episodes. |
|
| 31 |
-
| `test_anti_hivemind_protocol.py` | `cortex/` | 5 protocol steps fire in order; caps enforced (2 rounds, 1 cross-brain challenge, 1 Critic/brain/tick). |
|
| 32 |
-
| `test_collapse_detector.py` | `cortex/` | Metacognition flags when all brains recommend the same action. |
|
| 33 |
-
| `test_import_graph.py` | repo-wide | No `import server` under `cortex/**`; no `import cortex` under `server/**`; no `import server.simulator` under `training/**`. |
|
| 34 |
-
| `test_baselines_smoke.py` | `baselines/` | B1 / B2 / B3 each run one episode on `outbreak_easy`. |
|
| 35 |
-
| `test_training_smoke.py` | `training/` | `train_router.main()` runs one episode against a mocked env under 5 s. |
|
| 36 |
-
|
| 37 |
-
## Binding rules
|
| 38 |
-
|
| 39 |
-
- Every public API in a subsystem's CLAUDE.md has ≥ 1 test here.
|
| 40 |
-
- Coverage target: 80% per subsystem; 100% for `server/graders/` and `cortex/anti_hivemind.py`.
|
| 41 |
-
- No test may take > 10 s unless marked `@pytest.mark.slow` and gated behind `--runslow`.
|
| 42 |
-
- `test_import_graph.py` uses a fresh subprocess import, not `sys.modules` monkey-patching — the latter passes under contamination.
|
| 43 |
-
|
| 44 |
-
## Common failure modes
|
| 45 |
-
|
| 46 |
-
- Smoke test asserting on current-echo values — breaks when real env logic lands. Assert on shape, not value.
|
| 47 |
-
- Module-scope env instantiation in tests — slows collection and hides init errors until runtime.
|
| 48 |
-
- Tests that hit the HF Space without a skip guard — CI flakes on rate limits.
|
|
|
|
| 1 |
+
# tests/CLAUDE.md
|
| 2 |
+
|
| 3 |
+
Test surface per subsystem. Smoke bar first, boundary tests next, coverage last.
|
| 4 |
+
|
| 5 |
+
## Belongs here
|
| 6 |
+
|
| 7 |
+
- `conftest.py` — repo root on `sys.path` for bare-name imports.
|
| 8 |
+
- One test module per subsystem boundary (table below).
|
| 9 |
+
|
| 10 |
+
## Does not belong here
|
| 11 |
+
|
| 12 |
+
Helpers that mutate real graders, simulator state, or disk. Fixtures that hit the live HF Space — mock or `pytest.skip`.
|
| 13 |
+
|
| 14 |
+
## Run commands
|
| 15 |
+
|
| 16 |
+
```bash
|
| 17 |
+
uv run python -m pytest tests/ -v # all
|
| 18 |
+
uv run python -m pytest tests/test_smoke_env.py::test_reset_returns_valid_observation -v # one
|
| 19 |
+
uv run python -m pytest --cov tests/ # coverage
|
| 20 |
+
```
|
| 21 |
+
|
| 22 |
+
## Required tests — each maps to exactly one subsystem contract
|
| 23 |
+
|
| 24 |
+
| File | Scope | Asserts |
|
| 25 |
+
|---|---|---|
|
| 26 |
+
| `test_package_exports.py` | wire package | Root `__init__` re-exports `CrisisworldcortexAction/Observation/Env`. |
|
| 27 |
+
| `test_smoke_env.py` | `server/` env | `reset()` / `step()` return a valid `CrisisworldcortexObservation`. |
|
| 28 |
+
| `test_actions_round_trip.py` | `server/` env | 6 MVP outer actions round-trip; `public_communication` is rejected at runtime. |
|
| 29 |
+
| `test_reward_shape.py` | `server/graders/` | Every grader returns values in `[0.0, 1.0]`. |
|
| 30 |
+
| `test_reward_non_constancy.py` | `server/graders/` | Grader output varies across ≥ 2 synthetic episodes. |
|
| 31 |
+
| `test_anti_hivemind_protocol.py` | `cortex/` | 5 protocol steps fire in order; caps enforced (2 rounds, 1 cross-brain challenge, 1 Critic/brain/tick). |
|
| 32 |
+
| `test_collapse_detector.py` | `cortex/` | Metacognition flags when all brains recommend the same action. |
|
| 33 |
+
| `test_import_graph.py` | repo-wide | No `import server` under `cortex/**`; no `import cortex` under `server/**`; no `import server.simulator` under `training/**`. |
|
| 34 |
+
| `test_baselines_smoke.py` | `baselines/` | B1 / B2 / B3 each run one episode on `outbreak_easy`. |
|
| 35 |
+
| `test_training_smoke.py` | `training/` | `train_router.main()` runs one episode against a mocked env under 5 s. |
|
| 36 |
+
|
| 37 |
+
## Binding rules
|
| 38 |
+
|
| 39 |
+
- Every public API in a subsystem's CLAUDE.md has ≥ 1 test here.
|
| 40 |
+
- Coverage target: 80% per subsystem; 100% for `server/graders/` and `cortex/anti_hivemind.py`.
|
| 41 |
+
- No test may take > 10 s unless marked `@pytest.mark.slow` and gated behind `--runslow`.
|
| 42 |
+
- `test_import_graph.py` uses a fresh subprocess import, not `sys.modules` monkey-patching — the latter passes under contamination.
|
| 43 |
+
|
| 44 |
+
## Common failure modes
|
| 45 |
+
|
| 46 |
+
- Smoke test asserting on current-echo values — breaks when real env logic lands. Assert on shape, not value.
|
| 47 |
+
- Module-scope env instantiation in tests — slows collection and hides init errors until runtime.
|
| 48 |
+
- Tests that hit the HF Space without a skip guard — CI flakes on rate limits.
|
tests/_helpers/__init__.py
ADDED
|
@@ -0,0 +1,6 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Reusable test helpers for cortex/baselines tests.
|
| 2 |
+
|
| 3 |
+
Lives alongside ``tests/`` (not under ``cortex/``) so production code
|
| 4 |
+
never imports test doubles. Sessions 9-13 reuse the LLMClient stub
|
| 5 |
+
defined here.
|
| 6 |
+
"""
|
tests/_helpers/llm_stub.py
ADDED
|
@@ -0,0 +1,82 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""LLMClient-level test double.
|
| 2 |
+
|
| 3 |
+
Drop-in replacement for ``cortex.llm_client.LLMClient`` for tests that
|
| 4 |
+
exercise consumers of the client (subagents, brains, council). Differs
|
| 5 |
+
from the SDK-level stub in ``tests/test_llm_client.py``: that one
|
| 6 |
+
intercepts the OpenAI SDK; this one intercepts the ``LLMClient.chat``
|
| 7 |
+
surface directly, which is what subagents and harnesses see.
|
| 8 |
+
|
| 9 |
+
Reused by sessions 9-13 — do not bury role-specific test logic here.
|
| 10 |
+
"""
|
| 11 |
+
|
| 12 |
+
from __future__ import annotations
|
| 13 |
+
|
| 14 |
+
from dataclasses import dataclass, field
|
| 15 |
+
from typing import Dict, List, Optional
|
| 16 |
+
|
| 17 |
+
from cortex.llm_client import ChatMessage, ChatResponse
|
| 18 |
+
|
| 19 |
+
|
| 20 |
+
@dataclass
|
| 21 |
+
class RecordedCall:
|
| 22 |
+
"""One captured ``chat()`` invocation. Tests assert on these."""
|
| 23 |
+
|
| 24 |
+
caller_id: str
|
| 25 |
+
messages: List[ChatMessage]
|
| 26 |
+
max_tokens: Optional[int]
|
| 27 |
+
temperature: Optional[float]
|
| 28 |
+
|
| 29 |
+
|
| 30 |
+
@dataclass
|
| 31 |
+
class StubLLMClient:
|
| 32 |
+
"""Quacks like ``LLMClient`` for the ``chat()`` and counter surface.
|
| 33 |
+
|
| 34 |
+
Args:
|
| 35 |
+
scripted_responses: Yielded one per ``chat()`` call, in order.
|
| 36 |
+
Each entry is the response ``content`` string.
|
| 37 |
+
prompt_tokens_per_call: Fake prompt-token billing per call.
|
| 38 |
+
completion_tokens_per_call: Fake completion-token billing per call.
|
| 39 |
+
"""
|
| 40 |
+
|
| 41 |
+
scripted_responses: List[str]
|
| 42 |
+
prompt_tokens_per_call: int = 50
|
| 43 |
+
completion_tokens_per_call: int = 30
|
| 44 |
+
calls: List[RecordedCall] = field(default_factory=list)
|
| 45 |
+
_counters: Dict[str, int] = field(default_factory=dict)
|
| 46 |
+
|
| 47 |
+
def chat(
|
| 48 |
+
self,
|
| 49 |
+
caller_id: str,
|
| 50 |
+
messages: List[ChatMessage],
|
| 51 |
+
max_tokens: Optional[int] = None,
|
| 52 |
+
temperature: Optional[float] = None,
|
| 53 |
+
) -> ChatResponse:
|
| 54 |
+
if not self.scripted_responses:
|
| 55 |
+
raise RuntimeError(
|
| 56 |
+
f"StubLLMClient exhausted: caller_id={caller_id!r}; "
|
| 57 |
+
f"add more scripted_responses or assert call_count earlier."
|
| 58 |
+
)
|
| 59 |
+
content = self.scripted_responses.pop(0)
|
| 60 |
+
self.calls.append(
|
| 61 |
+
RecordedCall(
|
| 62 |
+
caller_id=caller_id,
|
| 63 |
+
messages=list(messages),
|
| 64 |
+
max_tokens=max_tokens,
|
| 65 |
+
temperature=temperature,
|
| 66 |
+
)
|
| 67 |
+
)
|
| 68 |
+
billed = self.prompt_tokens_per_call + self.completion_tokens_per_call
|
| 69 |
+
self._counters[caller_id] = self._counters.get(caller_id, 0) + billed
|
| 70 |
+
return ChatResponse(
|
| 71 |
+
content=content,
|
| 72 |
+
finish_reason="stop",
|
| 73 |
+
prompt_tokens=self.prompt_tokens_per_call,
|
| 74 |
+
completion_tokens=self.completion_tokens_per_call,
|
| 75 |
+
)
|
| 76 |
+
|
| 77 |
+
def tokens_used_for(self, caller_id: str) -> int:
|
| 78 |
+
return self._counters.get(caller_id, 0)
|
| 79 |
+
|
| 80 |
+
@property
|
| 81 |
+
def call_count(self) -> int:
|
| 82 |
+
return len(self.calls)
|
tests/test_baseline_b1.py
CHANGED
|
@@ -218,7 +218,13 @@ def test_b1_runs_episode_with_valid_json() -> None:
|
|
| 218 |
def test_b1_parse_failure_submits_synthetic_rejection() -> None:
|
| 219 |
"""When the LLM emits unparseable text, B1 submits a synthetic
|
| 220 |
PublicCommunication so the env rejects with accepted=False — landing
|
| 221 |
-
r_policy=0 in outer_reward
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 222 |
"""
|
| 223 |
env_inner = CrisisworldcortexEnvironment()
|
| 224 |
env = _InProcessEnvAdapter(env_inner)
|
|
@@ -233,36 +239,26 @@ def test_b1_parse_failure_submits_synthetic_rejection() -> None:
|
|
| 233 |
|
| 234 |
trajectory = agent.run_episode(task="outbreak_easy", seed=0, max_ticks=5)
|
| 235 |
|
| 236 |
-
#
|
| 237 |
-
assert trajectory["parse_failure_count"] ==
|
| 238 |
|
| 239 |
-
# B1 did NOT crash on parse failure —
|
| 240 |
-
#
|
| 241 |
-
|
| 242 |
-
assert trajectory["steps_taken"] >= 3, (
|
| 243 |
f"steps_taken={trajectory['steps_taken']!r} - parse failure "
|
| 244 |
-
f"
|
| 245 |
)
|
| 246 |
|
| 247 |
-
# The
|
| 248 |
-
#
|
| 249 |
log = env_inner._world_state.recent_action_log
|
| 250 |
-
assert len(log)
|
| 251 |
assert log[0].action.kind == "public_communication"
|
| 252 |
assert log[0].accepted is False, "parse-failure synthetic must be rejected by env"
|
| 253 |
-
assert log[1].action.kind == "public_communication"
|
| 254 |
-
assert log[1].accepted is False
|
| 255 |
-
|
| 256 |
-
# The third entry should be the first parsed NoOp.
|
| 257 |
-
assert log[2].action.kind == "no_op"
|
| 258 |
-
assert log[2].accepted is True
|
| 259 |
|
| 260 |
-
# B1's local trajectory carries the raw
|
| 261 |
assert trajectory["action_history"][0]["parse_failure"] is True
|
| 262 |
assert trajectory["action_history"][0]["raw_llm"] == "I cannot help with that."
|
| 263 |
-
assert trajectory["action_history"][1]["parse_failure"] is True
|
| 264 |
-
assert trajectory["action_history"][1]["raw_llm"] == "Sorry, no JSON."
|
| 265 |
-
assert trajectory["action_history"][2]["parse_failure"] is False
|
| 266 |
|
| 267 |
|
| 268 |
def test_b1_caller_id_format_short_colon_separated() -> None:
|
|
@@ -335,7 +331,10 @@ def test_b1_step_event_carries_rich_context() -> None:
|
|
| 335 |
from baselines.flat_agent import B1StepEvent
|
| 336 |
|
| 337 |
env = _InProcessEnvAdapter(CrisisworldcortexEnvironment())
|
| 338 |
-
|
|
|
|
|
|
|
|
|
|
| 339 |
agent = B1FlatAgent(env=env, llm=llm)
|
| 340 |
|
| 341 |
events: list[B1StepEvent] = []
|
|
@@ -348,18 +347,17 @@ def test_b1_step_event_carries_rich_context() -> None:
|
|
| 348 |
|
| 349 |
assert len(events) >= 2
|
| 350 |
|
| 351 |
-
# Tick 1:
|
| 352 |
-
# PublicCommunication marker; env returns accepted=False; reward in [0,1].
|
| 353 |
e1 = events[0]
|
| 354 |
assert e1.tick == 1
|
| 355 |
-
assert e1.parse_failure is
|
| 356 |
-
assert e1.error
|
| 357 |
-
assert e1.
|
| 358 |
-
assert e1.
|
| 359 |
-
assert
|
| 360 |
assert isinstance(e1.done, bool)
|
| 361 |
|
| 362 |
-
# Tick 2: clean parse. Submitted is NoOp. error must be None.
|
| 363 |
e2 = events[1]
|
| 364 |
assert e2.tick == 2
|
| 365 |
assert e2.parse_failure is False
|
|
|
|
| 218 |
def test_b1_parse_failure_submits_synthetic_rejection() -> None:
|
| 219 |
"""When the LLM emits unparseable text, B1 submits a synthetic
|
| 220 |
PublicCommunication so the env rejects with accepted=False — landing
|
| 221 |
+
r_policy=-1.0 in outer_reward (Phase-1 fix per design §19) and the
|
| 222 |
+
action log shows the rejection.
|
| 223 |
+
|
| 224 |
+
Phase-1 contract: parse-failure now TERMINATES the episode at the
|
| 225 |
+
rejection tick (state.terminal = "failure" → obs.done = True). So
|
| 226 |
+
only the first parse-failure lands; the second LLM response in the
|
| 227 |
+
stub queue never gets dispatched.
|
| 228 |
"""
|
| 229 |
env_inner = CrisisworldcortexEnvironment()
|
| 230 |
env = _InProcessEnvAdapter(env_inner)
|
|
|
|
| 239 |
|
| 240 |
trajectory = agent.run_episode(task="outbreak_easy", seed=0, max_ticks=5)
|
| 241 |
|
| 242 |
+
# Phase-1: parse-failure terminates → only the first marker lands.
|
| 243 |
+
assert trajectory["parse_failure_count"] == 1
|
| 244 |
|
| 245 |
+
# B1 did NOT crash on parse failure — exactly 1 tick ran (parse-failure
|
| 246 |
+
# marker submitted, env terminated episode).
|
| 247 |
+
assert trajectory["steps_taken"] == 1, (
|
|
|
|
| 248 |
f"steps_taken={trajectory['steps_taken']!r} - parse failure "
|
| 249 |
+
f"should terminate at tick 1 under the §19 contract"
|
| 250 |
)
|
| 251 |
|
| 252 |
+
# The action-log entry shows synthetic public_communication
|
| 253 |
+
# (parse-failure marker, honesty=0.0) — rejected by env.
|
| 254 |
log = env_inner._world_state.recent_action_log
|
| 255 |
+
assert len(log) == 1
|
| 256 |
assert log[0].action.kind == "public_communication"
|
| 257 |
assert log[0].accepted is False, "parse-failure synthetic must be rejected by env"
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 258 |
|
| 259 |
+
# B1's local trajectory carries the raw snippet for forensic use.
|
| 260 |
assert trajectory["action_history"][0]["parse_failure"] is True
|
| 261 |
assert trajectory["action_history"][0]["raw_llm"] == "I cannot help with that."
|
|
|
|
|
|
|
|
|
|
| 262 |
|
| 263 |
|
| 264 |
def test_b1_caller_id_format_short_colon_separated() -> None:
|
|
|
|
| 331 |
from baselines.flat_agent import B1StepEvent
|
| 332 |
|
| 333 |
env = _InProcessEnvAdapter(CrisisworldcortexEnvironment())
|
| 334 |
+
# Use a clean-parse first response; parse-failure now terminates the
|
| 335 |
+
# episode at tick 1 under the §19 contract, so we exercise the
|
| 336 |
+
# tick-1 + tick-2 sequence with both responses being valid JSON.
|
| 337 |
+
llm = _StubLLMClient(['{"kind": "no_op"}'] * 5)
|
| 338 |
agent = B1FlatAgent(env=env, llm=llm)
|
| 339 |
|
| 340 |
events: list[B1StepEvent] = []
|
|
|
|
| 347 |
|
| 348 |
assert len(events) >= 2
|
| 349 |
|
| 350 |
+
# Tick 1: clean parse. Submitted is NoOp. error must be None.
|
|
|
|
| 351 |
e1 = events[0]
|
| 352 |
assert e1.tick == 1
|
| 353 |
+
assert e1.parse_failure is False
|
| 354 |
+
assert e1.error is None
|
| 355 |
+
assert e1.action.kind == "no_op"
|
| 356 |
+
assert e1.raw_llm == '{"kind": "no_op"}'
|
| 357 |
+
assert -1.0 <= e1.reward <= 1.0 # Phase-1 range relaxed from [0,1].
|
| 358 |
assert isinstance(e1.done, bool)
|
| 359 |
|
| 360 |
+
# Tick 2: another clean parse. Submitted is NoOp. error must be None.
|
| 361 |
e2 = events[1]
|
| 362 |
assert e2.tick == 2
|
| 363 |
assert e2.parse_failure is False
|
tests/test_cortex_brain_executive.py
ADDED
|
@@ -0,0 +1,206 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Session 11 - Brain Executive aggregation tests.
|
| 2 |
+
|
| 3 |
+
Per Phase A docs/CORTEX_ARCHITECTURE.md Decisions 15-21 + M-FR-3
|
| 4 |
+
(partial evidence union; CandidatePlan and CriticReport schemas have
|
| 5 |
+
no evidence field).
|
| 6 |
+
"""
|
| 7 |
+
|
| 8 |
+
from __future__ import annotations
|
| 9 |
+
|
| 10 |
+
import pytest
|
| 11 |
+
|
| 12 |
+
from cortex.brains import aggregate_brain_outputs
|
| 13 |
+
from cortex.schemas import (
|
| 14 |
+
BeliefState,
|
| 15 |
+
CandidatePlan,
|
| 16 |
+
CriticReport,
|
| 17 |
+
EvidenceCitation,
|
| 18 |
+
Hypothesis,
|
| 19 |
+
PerceptionReport,
|
| 20 |
+
RegionBeliefEstimate,
|
| 21 |
+
)
|
| 22 |
+
from CrisisWorldCortex.models import (
|
| 23 |
+
DeployResource,
|
| 24 |
+
NoOp,
|
| 25 |
+
RestrictMovement,
|
| 26 |
+
)
|
| 27 |
+
|
| 28 |
+
|
| 29 |
+
def _belief(uncertainty: float = 0.4, evidence_count: int = 1) -> BeliefState:
|
| 30 |
+
return BeliefState(
|
| 31 |
+
brain="epidemiology",
|
| 32 |
+
latent_estimates={
|
| 33 |
+
"R1": RegionBeliefEstimate(
|
| 34 |
+
estimated_infection_rate=0.05,
|
| 35 |
+
estimated_r_effective=1.2,
|
| 36 |
+
estimated_compliance=0.85,
|
| 37 |
+
),
|
| 38 |
+
},
|
| 39 |
+
hypotheses=[Hypothesis(label="h1", weight=0.6, explanation="rising")],
|
| 40 |
+
uncertainty=uncertainty,
|
| 41 |
+
reducible_by_more_thought=0.3,
|
| 42 |
+
evidence=[
|
| 43 |
+
EvidenceCitation(source="telemetry", ref=f"R1.cases@{i}", excerpt=f"e{i}")
|
| 44 |
+
for i in range(evidence_count)
|
| 45 |
+
],
|
| 46 |
+
)
|
| 47 |
+
|
| 48 |
+
|
| 49 |
+
def _plan(action=None, confidence: float = 0.75, expected_value: float = 0.6) -> CandidatePlan:
|
| 50 |
+
if action is None:
|
| 51 |
+
action = DeployResource(region="R1", resource_type="test_kits", quantity=100)
|
| 52 |
+
return CandidatePlan(
|
| 53 |
+
action_sketch="Deploy 100 test_kits to R1",
|
| 54 |
+
expected_outer_action=action,
|
| 55 |
+
expected_value=expected_value,
|
| 56 |
+
cost=200.0,
|
| 57 |
+
assumptions=["kits available"],
|
| 58 |
+
falsifiers=["R1 cases drop without intervention"],
|
| 59 |
+
confidence=confidence,
|
| 60 |
+
)
|
| 61 |
+
|
| 62 |
+
|
| 63 |
+
def _critic(severity: float = 0.3) -> CriticReport:
|
| 64 |
+
return CriticReport(
|
| 65 |
+
brain="epidemiology",
|
| 66 |
+
target_plan_id="plan-0",
|
| 67 |
+
attacks=["limited reach"],
|
| 68 |
+
missing_considerations=[],
|
| 69 |
+
would_change_mind_if=[],
|
| 70 |
+
severity=severity,
|
| 71 |
+
)
|
| 72 |
+
|
| 73 |
+
|
| 74 |
+
def _perception(evidence_count: int = 1) -> PerceptionReport:
|
| 75 |
+
return PerceptionReport(
|
| 76 |
+
brain="epidemiology",
|
| 77 |
+
salient_signals=["R1 cases rising"],
|
| 78 |
+
anomalies=[],
|
| 79 |
+
confidence=0.7,
|
| 80 |
+
evidence=[
|
| 81 |
+
EvidenceCitation(source="telemetry", ref=f"R1.perception@{i}", excerpt=f"p{i}")
|
| 82 |
+
for i in range(evidence_count)
|
| 83 |
+
],
|
| 84 |
+
)
|
| 85 |
+
|
| 86 |
+
|
| 87 |
+
# T4
|
| 88 |
+
def test_brain_executive_aggregates_subagent_outputs() -> None:
|
| 89 |
+
rec = aggregate_brain_outputs(
|
| 90 |
+
brain_id="epidemiology",
|
| 91 |
+
perception=_perception(),
|
| 92 |
+
beliefs=[_belief()],
|
| 93 |
+
plans=[_plan()],
|
| 94 |
+
critics=[_critic()],
|
| 95 |
+
)
|
| 96 |
+
assert rec.brain == "epidemiology"
|
| 97 |
+
assert rec.top_action.kind == "deploy_resource"
|
| 98 |
+
assert rec.top_confidence > 0.0
|
| 99 |
+
assert rec.tokens_used == 0
|
| 100 |
+
|
| 101 |
+
|
| 102 |
+
# T5 -- Decision 16
|
| 103 |
+
def test_brain_executive_top_confidence_includes_uncertainty() -> None:
|
| 104 |
+
rec = aggregate_brain_outputs(
|
| 105 |
+
brain_id="epidemiology",
|
| 106 |
+
perception=_perception(),
|
| 107 |
+
beliefs=[_belief(uncertainty=0.4)],
|
| 108 |
+
plans=[_plan(confidence=0.75)],
|
| 109 |
+
critics=[_critic()],
|
| 110 |
+
)
|
| 111 |
+
# D16: top_confidence == confidence x (1 - uncertainty) == 0.75 x 0.6 == 0.45
|
| 112 |
+
assert rec.top_confidence == pytest.approx(0.45)
|
| 113 |
+
|
| 114 |
+
|
| 115 |
+
# T6 -- Decision 17
|
| 116 |
+
def test_brain_executive_minority_actions_excludes_top() -> None:
|
| 117 |
+
plan_a = _plan(
|
| 118 |
+
action=DeployResource(region="R1", resource_type="test_kits", quantity=100),
|
| 119 |
+
confidence=0.8,
|
| 120 |
+
expected_value=0.7,
|
| 121 |
+
)
|
| 122 |
+
plan_b = _plan(
|
| 123 |
+
action=RestrictMovement(region="R1", severity="moderate"),
|
| 124 |
+
confidence=0.5,
|
| 125 |
+
expected_value=0.4,
|
| 126 |
+
)
|
| 127 |
+
# plan_a wins: 0.8 * 0.7 = 0.56 > 0.5 * 0.4 = 0.20
|
| 128 |
+
|
| 129 |
+
rec = aggregate_brain_outputs(
|
| 130 |
+
brain_id="epidemiology",
|
| 131 |
+
perception=_perception(),
|
| 132 |
+
beliefs=[_belief(), _belief(uncertainty=0.5)],
|
| 133 |
+
plans=[plan_a, plan_b],
|
| 134 |
+
critics=[_critic(), _critic()],
|
| 135 |
+
)
|
| 136 |
+
|
| 137 |
+
assert rec.top_action.kind == "deploy_resource"
|
| 138 |
+
assert len(rec.minority_actions) == 1
|
| 139 |
+
assert rec.minority_actions[0].kind == "restrict_movement"
|
| 140 |
+
|
| 141 |
+
|
| 142 |
+
# T7 -- Decision 20 + M-FR-3
|
| 143 |
+
def test_brain_executive_evidence_union() -> None:
|
| 144 |
+
perception = _perception(evidence_count=1)
|
| 145 |
+
belief = _belief(evidence_count=2)
|
| 146 |
+
|
| 147 |
+
rec = aggregate_brain_outputs(
|
| 148 |
+
brain_id="epidemiology",
|
| 149 |
+
perception=perception,
|
| 150 |
+
beliefs=[belief],
|
| 151 |
+
plans=[_plan()],
|
| 152 |
+
critics=[_critic()],
|
| 153 |
+
)
|
| 154 |
+
|
| 155 |
+
# M-FR-3: union of perception.evidence + belief.evidence (3 total)
|
| 156 |
+
assert len(rec.evidence) == 3
|
| 157 |
+
assert rec.evidence[0].ref.startswith("R1.perception")
|
| 158 |
+
assert rec.evidence[1].ref.startswith("R1.cases")
|
| 159 |
+
|
| 160 |
+
|
| 161 |
+
# T8 -- empty fallback
|
| 162 |
+
def test_brain_executive_handles_empty_subagent_outputs() -> None:
|
| 163 |
+
empty_belief = BeliefState(
|
| 164 |
+
brain="epidemiology",
|
| 165 |
+
latent_estimates={},
|
| 166 |
+
hypotheses=[],
|
| 167 |
+
uncertainty=1.0,
|
| 168 |
+
reducible_by_more_thought=0.0,
|
| 169 |
+
evidence=[],
|
| 170 |
+
)
|
| 171 |
+
empty_plan = CandidatePlan(
|
| 172 |
+
action_sketch="(empty)",
|
| 173 |
+
expected_outer_action=NoOp(),
|
| 174 |
+
expected_value=0.0,
|
| 175 |
+
cost=0.0,
|
| 176 |
+
assumptions=[],
|
| 177 |
+
falsifiers=[],
|
| 178 |
+
confidence=0.0,
|
| 179 |
+
)
|
| 180 |
+
empty_critic = CriticReport(
|
| 181 |
+
brain="epidemiology",
|
| 182 |
+
target_plan_id="",
|
| 183 |
+
attacks=[],
|
| 184 |
+
missing_considerations=[],
|
| 185 |
+
would_change_mind_if=[],
|
| 186 |
+
severity=0.0,
|
| 187 |
+
)
|
| 188 |
+
empty_perception = PerceptionReport(
|
| 189 |
+
brain="epidemiology",
|
| 190 |
+
salient_signals=[],
|
| 191 |
+
anomalies=[],
|
| 192 |
+
confidence=0.0,
|
| 193 |
+
evidence=[],
|
| 194 |
+
)
|
| 195 |
+
|
| 196 |
+
rec = aggregate_brain_outputs(
|
| 197 |
+
brain_id="epidemiology",
|
| 198 |
+
perception=empty_perception,
|
| 199 |
+
beliefs=[empty_belief],
|
| 200 |
+
plans=[empty_plan],
|
| 201 |
+
critics=[empty_critic],
|
| 202 |
+
)
|
| 203 |
+
|
| 204 |
+
assert rec.top_action.kind == "no_op"
|
| 205 |
+
assert rec.top_confidence == 0.0
|
| 206 |
+
assert rec.uncertainty == 1.0
|
tests/test_cortex_brain_smoke.py
ADDED
|
@@ -0,0 +1,142 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Session 11 - Single-brain end-to-end smoke + no-LLM-in-Python-layers test.
|
| 2 |
+
|
| 3 |
+
T9 asserts Brain.compute_perception and Brain.compute_lens are pure
|
| 4 |
+
Python (zero LLM calls). T10 is the integration smoke gate per Phase A
|
| 5 |
+
section 10: a single brain runs end-to-end on a real observation,
|
| 6 |
+
returns a BrainRecommendation. Three LLM calls in canonical
|
| 7 |
+
WorldModeler -> Planner -> Critic order with locked caller_id format.
|
| 8 |
+
"""
|
| 9 |
+
|
| 10 |
+
from __future__ import annotations
|
| 11 |
+
|
| 12 |
+
import json
|
| 13 |
+
|
| 14 |
+
from cortex.brains import EpiBrain
|
| 15 |
+
from cortex.schemas import BrainRecommendation
|
| 16 |
+
from CrisisWorldCortex.models import (
|
| 17 |
+
CrisisworldcortexObservation,
|
| 18 |
+
RegionTelemetry,
|
| 19 |
+
ResourceInventory,
|
| 20 |
+
)
|
| 21 |
+
from tests._helpers.llm_stub import StubLLMClient
|
| 22 |
+
|
| 23 |
+
|
| 24 |
+
def _make_obs() -> CrisisworldcortexObservation:
|
| 25 |
+
return CrisisworldcortexObservation(
|
| 26 |
+
regions=[
|
| 27 |
+
RegionTelemetry(
|
| 28 |
+
region=f"R{i + 1}",
|
| 29 |
+
reported_cases_d_ago=5 if i == 0 else 1,
|
| 30 |
+
hospital_load=0.3 if i == 0 else 0.1,
|
| 31 |
+
compliance_proxy=0.85,
|
| 32 |
+
)
|
| 33 |
+
for i in range(4)
|
| 34 |
+
],
|
| 35 |
+
resources=ResourceInventory(
|
| 36 |
+
test_kits=1000,
|
| 37 |
+
hospital_beds_free=500,
|
| 38 |
+
mobile_units=20,
|
| 39 |
+
vaccine_doses=2000,
|
| 40 |
+
),
|
| 41 |
+
active_restrictions=[],
|
| 42 |
+
legal_constraints=[],
|
| 43 |
+
tick=3,
|
| 44 |
+
ticks_remaining=9,
|
| 45 |
+
cognition_budget_remaining=5200,
|
| 46 |
+
recent_action_log=[],
|
| 47 |
+
)
|
| 48 |
+
|
| 49 |
+
|
| 50 |
+
_VALID_BELIEF = json.dumps(
|
| 51 |
+
{
|
| 52 |
+
"brain": "epidemiology",
|
| 53 |
+
"latent_estimates": {
|
| 54 |
+
"R1": {
|
| 55 |
+
"estimated_infection_rate": 0.05,
|
| 56 |
+
"estimated_r_effective": 1.2,
|
| 57 |
+
"estimated_compliance": 0.85,
|
| 58 |
+
"confidence_intervals": {},
|
| 59 |
+
}
|
| 60 |
+
},
|
| 61 |
+
"hypotheses": [{"label": "rising", "weight": 0.6, "explanation": "R1 cases up"}],
|
| 62 |
+
"uncertainty": 0.4,
|
| 63 |
+
"reducible_by_more_thought": 0.3,
|
| 64 |
+
"evidence": [
|
| 65 |
+
{"source": "telemetry", "ref": "R1.cases", "excerpt": "5"},
|
| 66 |
+
{"source": "policy", "ref": "R1.restriction", "excerpt": "none"},
|
| 67 |
+
],
|
| 68 |
+
}
|
| 69 |
+
)
|
| 70 |
+
|
| 71 |
+
_VALID_PLAN = json.dumps(
|
| 72 |
+
{
|
| 73 |
+
"action_sketch": "Deploy 100 test_kits to R1",
|
| 74 |
+
"expected_outer_action": {
|
| 75 |
+
"kind": "deploy_resource",
|
| 76 |
+
"region": "R1",
|
| 77 |
+
"resource_type": "test_kits",
|
| 78 |
+
"quantity": 100,
|
| 79 |
+
},
|
| 80 |
+
"expected_value": 0.6,
|
| 81 |
+
"cost": 200.0,
|
| 82 |
+
"assumptions": ["kits inventory > 100"],
|
| 83 |
+
"falsifiers": ["R1 cases drop without intervention"],
|
| 84 |
+
"confidence": 0.75,
|
| 85 |
+
}
|
| 86 |
+
)
|
| 87 |
+
|
| 88 |
+
_VALID_CRITIC = json.dumps(
|
| 89 |
+
{
|
| 90 |
+
"brain": "epidemiology",
|
| 91 |
+
"target_plan_id": "plan-0",
|
| 92 |
+
"attacks": ["ignores R3 hospital saturation"],
|
| 93 |
+
"missing_considerations": [],
|
| 94 |
+
"would_change_mind_if": [],
|
| 95 |
+
"severity": 0.3,
|
| 96 |
+
}
|
| 97 |
+
)
|
| 98 |
+
|
| 99 |
+
|
| 100 |
+
# T9
|
| 101 |
+
def test_brain_runs_zero_llm_calls_in_python_layers() -> None:
|
| 102 |
+
"""Perception and Lens are pure Python; no LLMClient invocation."""
|
| 103 |
+
stub = StubLLMClient(scripted_responses=[]) # any chat() would raise
|
| 104 |
+
brain = EpiBrain(stub)
|
| 105 |
+
obs = _make_obs()
|
| 106 |
+
|
| 107 |
+
perception = brain.compute_perception(obs)
|
| 108 |
+
lensed = brain.compute_lens(obs, last_reward=0.0)
|
| 109 |
+
|
| 110 |
+
assert stub.call_count == 0
|
| 111 |
+
assert perception.brain == "epidemiology"
|
| 112 |
+
assert lensed.brain == "epidemiology"
|
| 113 |
+
|
| 114 |
+
|
| 115 |
+
# T10 -- integration smoke gate (Phase A section 10)
|
| 116 |
+
def test_brain_smoke_one_tick_three_llm_calls_in_order() -> None:
|
| 117 |
+
"""Full round-1 tick: WorldModeler -> Planner -> Critic, in that order."""
|
| 118 |
+
stub = StubLLMClient(scripted_responses=[_VALID_BELIEF, _VALID_PLAN, _VALID_CRITIC])
|
| 119 |
+
brain = EpiBrain(stub)
|
| 120 |
+
obs = _make_obs()
|
| 121 |
+
|
| 122 |
+
rec = brain.run_tick(obs, last_reward=0.0, tick=3)
|
| 123 |
+
|
| 124 |
+
assert isinstance(rec, BrainRecommendation)
|
| 125 |
+
assert rec.brain == "epidemiology"
|
| 126 |
+
assert rec.top_action.kind == "deploy_resource"
|
| 127 |
+
assert stub.call_count == 3, "exactly 3 LLM calls per Phase A 'WM + Planner + Critic'"
|
| 128 |
+
|
| 129 |
+
# Order pin per user adjustment: WM (s0) -> Planner (s1) -> Critic (s2)
|
| 130 |
+
assert stub.calls[0].caller_id.endswith(":world_modeler:t3:r1:s0"), (
|
| 131 |
+
f"first call must be WorldModeler, got {stub.calls[0].caller_id!r}"
|
| 132 |
+
)
|
| 133 |
+
assert stub.calls[1].caller_id.endswith(":planner:t3:r1:s1"), (
|
| 134 |
+
f"second call must be Planner, got {stub.calls[1].caller_id!r}"
|
| 135 |
+
)
|
| 136 |
+
assert stub.calls[2].caller_id.endswith(":critic:t3:r1:s2"), (
|
| 137 |
+
f"third call must be Critic, got {stub.calls[2].caller_id!r}"
|
| 138 |
+
)
|
| 139 |
+
|
| 140 |
+
# Brain prefix locks
|
| 141 |
+
for call in stub.calls:
|
| 142 |
+
assert call.caller_id.startswith("cortex:epidemiology:")
|
tests/test_cortex_lenses.py
ADDED
|
@@ -0,0 +1,238 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Session 10 - Cortex lens tests.
|
| 2 |
+
|
| 3 |
+
Per Phase A docs/CORTEX_ARCHITECTURE.md Decisions 9-14 + §2 A1 and the
|
| 4 |
+
user's Session 10 proposal acceptance with 8 tests + the M-FR-4 rename
|
| 5 |
+
(epi_pressure) and T7 tightening (non-bool float).
|
| 6 |
+
"""
|
| 7 |
+
|
| 8 |
+
from __future__ import annotations
|
| 9 |
+
|
| 10 |
+
from typing import Iterable
|
| 11 |
+
|
| 12 |
+
import pytest
|
| 13 |
+
|
| 14 |
+
from cortex.lenses import lens_for
|
| 15 |
+
from cortex.schemas import BrainLensedObservation
|
| 16 |
+
from CrisisWorldCortex.models import (
|
| 17 |
+
CrisisworldcortexObservation,
|
| 18 |
+
Escalate,
|
| 19 |
+
ExecutedAction,
|
| 20 |
+
LegalConstraint,
|
| 21 |
+
NoOp,
|
| 22 |
+
RegionTelemetry,
|
| 23 |
+
ResourceInventory,
|
| 24 |
+
Restriction,
|
| 25 |
+
)
|
| 26 |
+
|
| 27 |
+
# ============================================================================
|
| 28 |
+
# Test fixtures
|
| 29 |
+
# ============================================================================
|
| 30 |
+
|
| 31 |
+
|
| 32 |
+
def _make_obs(
|
| 33 |
+
cases_per_region: Iterable[int] = (5, 1, 1, 1),
|
| 34 |
+
hospital_loads: Iterable[float] = (0.3, 0.1, 0.1, 0.1),
|
| 35 |
+
compliance_proxies: Iterable[float] = (0.85, 0.95, 0.95, 0.95),
|
| 36 |
+
test_kits: int = 1000,
|
| 37 |
+
hospital_beds_free: int = 500,
|
| 38 |
+
mobile_units: int = 20,
|
| 39 |
+
vaccine_doses: int = 2000,
|
| 40 |
+
restrictions: Iterable[Restriction] = (),
|
| 41 |
+
legal_constraints: Iterable[LegalConstraint] = (),
|
| 42 |
+
recent_action_log: Iterable[ExecutedAction] = (),
|
| 43 |
+
) -> CrisisworldcortexObservation:
|
| 44 |
+
cases = list(cases_per_region)
|
| 45 |
+
loads = list(hospital_loads)
|
| 46 |
+
comps = list(compliance_proxies)
|
| 47 |
+
return CrisisworldcortexObservation(
|
| 48 |
+
regions=[
|
| 49 |
+
RegionTelemetry(
|
| 50 |
+
region=f"R{i + 1}",
|
| 51 |
+
reported_cases_d_ago=cases[i],
|
| 52 |
+
hospital_load=loads[i],
|
| 53 |
+
compliance_proxy=comps[i],
|
| 54 |
+
)
|
| 55 |
+
for i in range(4)
|
| 56 |
+
],
|
| 57 |
+
resources=ResourceInventory(
|
| 58 |
+
test_kits=test_kits,
|
| 59 |
+
hospital_beds_free=hospital_beds_free,
|
| 60 |
+
mobile_units=mobile_units,
|
| 61 |
+
vaccine_doses=vaccine_doses,
|
| 62 |
+
),
|
| 63 |
+
active_restrictions=list(restrictions),
|
| 64 |
+
legal_constraints=list(legal_constraints),
|
| 65 |
+
tick=3,
|
| 66 |
+
ticks_remaining=9,
|
| 67 |
+
cognition_budget_remaining=5200,
|
| 68 |
+
recent_action_log=list(recent_action_log),
|
| 69 |
+
)
|
| 70 |
+
|
| 71 |
+
|
| 72 |
+
# ============================================================================
|
| 73 |
+
# T1 - Epi lens emphasizes telemetry; uses epi_pressure (M-FR-4 rename)
|
| 74 |
+
# ============================================================================
|
| 75 |
+
|
| 76 |
+
|
| 77 |
+
def test_epi_lens_emphasizes_telemetry() -> None:
|
| 78 |
+
obs = _make_obs()
|
| 79 |
+
lensed = lens_for("epidemiology", obs, last_reward=0.5)
|
| 80 |
+
|
| 81 |
+
assert isinstance(lensed, BrainLensedObservation)
|
| 82 |
+
assert lensed.brain == "epidemiology"
|
| 83 |
+
assert lensed.last_reward == 0.5
|
| 84 |
+
|
| 85 |
+
keys = set(lensed.derived_features.keys())
|
| 86 |
+
assert {"epi_pressure", "worst_region_infection", "transmission_rate_trend"} <= keys
|
| 87 |
+
|
| 88 |
+
# M-FR-2: trend is 0.0 in MVP (no history available in single-obs lens)
|
| 89 |
+
assert lensed.derived_features["transmission_rate_trend"] == 0.0
|
| 90 |
+
|
| 91 |
+
assert "regions[*].reported_cases_d_ago" in lensed.salient_field_ids
|
| 92 |
+
assert "regions[*].hospital_load" in lensed.salient_field_ids
|
| 93 |
+
|
| 94 |
+
|
| 95 |
+
# ============================================================================
|
| 96 |
+
# T2 - Logistics lens emphasizes resources
|
| 97 |
+
# ============================================================================
|
| 98 |
+
|
| 99 |
+
|
| 100 |
+
def test_logistics_lens_emphasizes_resources() -> None:
|
| 101 |
+
obs = _make_obs(test_kits=100, hospital_beds_free=50, mobile_units=10, vaccine_doses=200)
|
| 102 |
+
lensed = lens_for("logistics", obs, last_reward=0.5)
|
| 103 |
+
|
| 104 |
+
assert lensed.brain == "logistics"
|
| 105 |
+
keys = set(lensed.derived_features.keys())
|
| 106 |
+
expected_keys = {
|
| 107 |
+
"total_inventory",
|
| 108 |
+
"hospital_load_max",
|
| 109 |
+
"deployment_feasibility_R1",
|
| 110 |
+
"deployment_feasibility_R2",
|
| 111 |
+
"deployment_feasibility_R3",
|
| 112 |
+
"deployment_feasibility_R4",
|
| 113 |
+
}
|
| 114 |
+
assert expected_keys <= keys
|
| 115 |
+
|
| 116 |
+
# 100 + 50 + 10 + 200 = 360
|
| 117 |
+
assert lensed.derived_features["total_inventory"] == 360.0
|
| 118 |
+
|
| 119 |
+
assert "resources.test_kits" in lensed.salient_field_ids
|
| 120 |
+
|
| 121 |
+
|
| 122 |
+
# ============================================================================
|
| 123 |
+
# T3 - Governance lens emphasizes legal
|
| 124 |
+
# ============================================================================
|
| 125 |
+
|
| 126 |
+
|
| 127 |
+
def test_governance_lens_emphasizes_legal() -> None:
|
| 128 |
+
obs = _make_obs(
|
| 129 |
+
restrictions=[Restriction(region="R1", severity="moderate", ticks_remaining=3)],
|
| 130 |
+
legal_constraints=[
|
| 131 |
+
LegalConstraint(rule_id="L1", blocked_action="restrict_movement.strict")
|
| 132 |
+
],
|
| 133 |
+
)
|
| 134 |
+
lensed = lens_for("governance", obs, last_reward=0.5)
|
| 135 |
+
|
| 136 |
+
assert lensed.brain == "governance"
|
| 137 |
+
keys = set(lensed.derived_features.keys())
|
| 138 |
+
assert {
|
| 139 |
+
"escalation_unlocked_strict",
|
| 140 |
+
"legal_constraints_count",
|
| 141 |
+
"restrictions_active_count",
|
| 142 |
+
} <= keys
|
| 143 |
+
|
| 144 |
+
assert lensed.derived_features["legal_constraints_count"] == 1.0
|
| 145 |
+
assert lensed.derived_features["restrictions_active_count"] == 1.0
|
| 146 |
+
assert "active_restrictions[*]" in lensed.salient_field_ids
|
| 147 |
+
|
| 148 |
+
|
| 149 |
+
# ============================================================================
|
| 150 |
+
# T4 - Lens does NOT strip raw_obs (D13)
|
| 151 |
+
# ============================================================================
|
| 152 |
+
|
| 153 |
+
|
| 154 |
+
def test_lens_does_not_strip_raw_obs() -> None:
|
| 155 |
+
obs = _make_obs(restrictions=[Restriction(region="R1", severity="moderate", ticks_remaining=3)])
|
| 156 |
+
|
| 157 |
+
for brain in ("epidemiology", "logistics", "governance"):
|
| 158 |
+
lensed = lens_for(brain, obs, last_reward=0.0)
|
| 159 |
+
# Pydantic deep-equality on the full observation
|
| 160 |
+
assert lensed.raw_obs == obs
|
| 161 |
+
|
| 162 |
+
|
| 163 |
+
# ============================================================================
|
| 164 |
+
# T5 - V2 brain ids raise KeyError (Decision 9, post-review)
|
| 165 |
+
# ============================================================================
|
| 166 |
+
|
| 167 |
+
|
| 168 |
+
def test_lens_for_v2_brain_raises_key_error() -> None:
|
| 169 |
+
obs = _make_obs()
|
| 170 |
+
|
| 171 |
+
for v2_brain in ("communications", "equity"):
|
| 172 |
+
with pytest.raises(KeyError):
|
| 173 |
+
lens_for(v2_brain, obs, last_reward=0.0)
|
| 174 |
+
|
| 175 |
+
with pytest.raises(KeyError):
|
| 176 |
+
lens_for("not_a_brain", obs, last_reward=0.0)
|
| 177 |
+
|
| 178 |
+
|
| 179 |
+
# ============================================================================
|
| 180 |
+
# T7 - All derived_features values are non-bool floats (D14 + tightened)
|
| 181 |
+
# ============================================================================
|
| 182 |
+
|
| 183 |
+
|
| 184 |
+
def test_lens_derived_features_all_floats() -> None:
|
| 185 |
+
obs = _make_obs(
|
| 186 |
+
restrictions=[Restriction(region="R1", severity="moderate", ticks_remaining=3)],
|
| 187 |
+
legal_constraints=[
|
| 188 |
+
LegalConstraint(rule_id="L1", blocked_action="restrict_movement.strict")
|
| 189 |
+
],
|
| 190 |
+
)
|
| 191 |
+
|
| 192 |
+
for brain in ("epidemiology", "logistics", "governance"):
|
| 193 |
+
lensed = lens_for(brain, obs, last_reward=0.0)
|
| 194 |
+
for key, value in lensed.derived_features.items():
|
| 195 |
+
assert isinstance(value, float) and not isinstance(value, bool), (
|
| 196 |
+
f"derived_features[{key!r}] = {value!r} ({type(value).__name__}) "
|
| 197 |
+
f"is not a non-bool float"
|
| 198 |
+
)
|
| 199 |
+
|
| 200 |
+
|
| 201 |
+
# ============================================================================
|
| 202 |
+
# T8 - Governance lens detects accepted escalate(national)
|
| 203 |
+
# ============================================================================
|
| 204 |
+
|
| 205 |
+
|
| 206 |
+
def test_governance_lens_detects_escalation_unlocked() -> None:
|
| 207 |
+
obs_with_accepted = _make_obs(
|
| 208 |
+
recent_action_log=[
|
| 209 |
+
ExecutedAction(tick=1, action=NoOp(), accepted=True),
|
| 210 |
+
ExecutedAction(tick=2, action=Escalate(to_authority="national"), accepted=True),
|
| 211 |
+
]
|
| 212 |
+
)
|
| 213 |
+
lensed = lens_for("governance", obs_with_accepted, last_reward=0.0)
|
| 214 |
+
assert lensed.derived_features["escalation_unlocked_strict"] == 1.0
|
| 215 |
+
|
| 216 |
+
obs_without = _make_obs(
|
| 217 |
+
recent_action_log=[ExecutedAction(tick=1, action=NoOp(), accepted=True)]
|
| 218 |
+
)
|
| 219 |
+
lensed_w = lens_for("governance", obs_without, last_reward=0.0)
|
| 220 |
+
assert lensed_w.derived_features["escalation_unlocked_strict"] == 0.0
|
| 221 |
+
|
| 222 |
+
# Rejected escalate must NOT count
|
| 223 |
+
obs_rejected = _make_obs(
|
| 224 |
+
recent_action_log=[
|
| 225 |
+
ExecutedAction(tick=1, action=Escalate(to_authority="national"), accepted=False),
|
| 226 |
+
]
|
| 227 |
+
)
|
| 228 |
+
lensed_r = lens_for("governance", obs_rejected, last_reward=0.0)
|
| 229 |
+
assert lensed_r.derived_features["escalation_unlocked_strict"] == 0.0
|
| 230 |
+
|
| 231 |
+
# Accepted escalate(regional) must NOT count -- only "national" unlocks strict
|
| 232 |
+
obs_regional = _make_obs(
|
| 233 |
+
recent_action_log=[
|
| 234 |
+
ExecutedAction(tick=1, action=Escalate(to_authority="regional"), accepted=True),
|
| 235 |
+
]
|
| 236 |
+
)
|
| 237 |
+
lensed_re = lens_for("governance", obs_regional, last_reward=0.0)
|
| 238 |
+
assert lensed_re.derived_features["escalation_unlocked_strict"] == 0.0
|
tests/test_cortex_perception.py
ADDED
|
@@ -0,0 +1,116 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Session 11 - Perception subagent tests (deterministic Python; no LLM).
|
| 2 |
+
|
| 3 |
+
Per cortex/CLAUDE.md binding (Perception is pure Python) and Phase A
|
| 4 |
+
Decisions 9 (V2 KeyError) + 63 (salient_signals cap at 5).
|
| 5 |
+
"""
|
| 6 |
+
|
| 7 |
+
from __future__ import annotations
|
| 8 |
+
|
| 9 |
+
from typing import Iterable
|
| 10 |
+
|
| 11 |
+
import pytest
|
| 12 |
+
|
| 13 |
+
from cortex.schemas import PerceptionReport
|
| 14 |
+
from cortex.subagents import perception_for
|
| 15 |
+
from CrisisWorldCortex.models import (
|
| 16 |
+
CrisisworldcortexObservation,
|
| 17 |
+
LegalConstraint,
|
| 18 |
+
RegionTelemetry,
|
| 19 |
+
ResourceInventory,
|
| 20 |
+
Restriction,
|
| 21 |
+
)
|
| 22 |
+
|
| 23 |
+
|
| 24 |
+
def _make_obs(
|
| 25 |
+
cases_per_region: Iterable[int] = (5, 1, 0, 0),
|
| 26 |
+
hospital_loads: Iterable[float] = (0.7, 0.1, 0.1, 0.1),
|
| 27 |
+
test_kits: int = 200,
|
| 28 |
+
hospital_beds_free: int = 50,
|
| 29 |
+
mobile_units: int = 2,
|
| 30 |
+
vaccine_doses: int = 300,
|
| 31 |
+
restrictions: Iterable[Restriction] = (),
|
| 32 |
+
legal_constraints: Iterable[LegalConstraint] = (),
|
| 33 |
+
) -> CrisisworldcortexObservation:
|
| 34 |
+
cases = list(cases_per_region)
|
| 35 |
+
loads = list(hospital_loads)
|
| 36 |
+
return CrisisworldcortexObservation(
|
| 37 |
+
regions=[
|
| 38 |
+
RegionTelemetry(
|
| 39 |
+
region=f"R{i + 1}",
|
| 40 |
+
reported_cases_d_ago=cases[i],
|
| 41 |
+
hospital_load=loads[i],
|
| 42 |
+
compliance_proxy=0.85,
|
| 43 |
+
)
|
| 44 |
+
for i in range(4)
|
| 45 |
+
],
|
| 46 |
+
resources=ResourceInventory(
|
| 47 |
+
test_kits=test_kits,
|
| 48 |
+
hospital_beds_free=hospital_beds_free,
|
| 49 |
+
mobile_units=mobile_units,
|
| 50 |
+
vaccine_doses=vaccine_doses,
|
| 51 |
+
),
|
| 52 |
+
active_restrictions=list(restrictions),
|
| 53 |
+
legal_constraints=list(legal_constraints),
|
| 54 |
+
tick=3,
|
| 55 |
+
ticks_remaining=9,
|
| 56 |
+
cognition_budget_remaining=5200,
|
| 57 |
+
recent_action_log=[],
|
| 58 |
+
)
|
| 59 |
+
|
| 60 |
+
|
| 61 |
+
# T1
|
| 62 |
+
def test_perception_runs_without_llm_call() -> None:
|
| 63 |
+
"""Perception is pure Python; the function does not take an LLMClient."""
|
| 64 |
+
obs = _make_obs()
|
| 65 |
+
for brain in ("epidemiology", "logistics", "governance"):
|
| 66 |
+
report = perception_for(brain, obs)
|
| 67 |
+
assert isinstance(report, PerceptionReport)
|
| 68 |
+
assert report.brain == brain
|
| 69 |
+
assert isinstance(report.confidence, float)
|
| 70 |
+
assert 0.0 <= report.confidence <= 1.0
|
| 71 |
+
|
| 72 |
+
|
| 73 |
+
# T2
|
| 74 |
+
def test_perception_for_v2_brain_raises_key_error() -> None:
|
| 75 |
+
obs = _make_obs()
|
| 76 |
+
for v2_brain in ("communications", "equity"):
|
| 77 |
+
with pytest.raises(KeyError):
|
| 78 |
+
perception_for(v2_brain, obs)
|
| 79 |
+
with pytest.raises(KeyError):
|
| 80 |
+
perception_for("not_a_brain", obs)
|
| 81 |
+
|
| 82 |
+
|
| 83 |
+
# T3
|
| 84 |
+
@pytest.mark.parametrize("brain", ["epidemiology", "logistics", "governance"])
|
| 85 |
+
def test_perception_brain_specific_signals(brain: str) -> None:
|
| 86 |
+
obs = _make_obs(
|
| 87 |
+
cases_per_region=(20, 1, 0, 0),
|
| 88 |
+
hospital_loads=(0.7, 0.1, 0.1, 0.1),
|
| 89 |
+
test_kits=100, # below threshold (300)
|
| 90 |
+
hospital_beds_free=50, # below threshold (100)
|
| 91 |
+
mobile_units=2, # below threshold (5)
|
| 92 |
+
vaccine_doses=200, # below threshold (500)
|
| 93 |
+
restrictions=[
|
| 94 |
+
Restriction(region="R1", severity="moderate", ticks_remaining=3),
|
| 95 |
+
],
|
| 96 |
+
legal_constraints=[
|
| 97 |
+
LegalConstraint(rule_id="L1", blocked_action="restrict_movement.strict"),
|
| 98 |
+
],
|
| 99 |
+
)
|
| 100 |
+
report = perception_for(brain, obs)
|
| 101 |
+
|
| 102 |
+
assert report.brain == brain
|
| 103 |
+
# Decision 63 / OQ-2 cap: at most 5 entries
|
| 104 |
+
assert len(report.salient_signals) <= 5
|
| 105 |
+
|
| 106 |
+
if brain == "epidemiology":
|
| 107 |
+
assert any("R1" in s for s in report.salient_signals), (
|
| 108 |
+
f"epi salient_signals should reference R1, got {report.salient_signals}"
|
| 109 |
+
)
|
| 110 |
+
elif brain == "logistics":
|
| 111 |
+
joined = " ".join(report.salient_signals).lower()
|
| 112 |
+
assert "kits" in joined or "mobile" in joined or "vaccine" in joined or "beds" in joined
|
| 113 |
+
elif brain == "governance":
|
| 114 |
+
assert any("R1" in s and "moderate" in s.lower() for s in report.salient_signals), (
|
| 115 |
+
f"governance salient_signals should mention R1 moderate, got {report.salient_signals}"
|
| 116 |
+
)
|
tests/test_cortex_subagents.py
ADDED
|
@@ -0,0 +1,387 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Session 9 - Cortex subagent tests (WorldModeler, Planner, Critic).
|
| 2 |
+
|
| 3 |
+
Per Phase A docs/CORTEX_ARCHITECTURE.md Decisions 1-8 + 62 and the user's
|
| 4 |
+
proposal acceptance with 11 tests total. RED-tests-first.
|
| 5 |
+
"""
|
| 6 |
+
|
| 7 |
+
from __future__ import annotations
|
| 8 |
+
|
| 9 |
+
import json
|
| 10 |
+
import re
|
| 11 |
+
from typing import Any, Dict, Optional
|
| 12 |
+
|
| 13 |
+
import pytest
|
| 14 |
+
|
| 15 |
+
from cortex.schemas import (
|
| 16 |
+
BeliefState,
|
| 17 |
+
CandidatePlan,
|
| 18 |
+
CriticReport,
|
| 19 |
+
EvidenceCitation,
|
| 20 |
+
PerceptionReport,
|
| 21 |
+
SubagentInput,
|
| 22 |
+
)
|
| 23 |
+
from cortex.subagents import (
|
| 24 |
+
PROMPTS_DIR,
|
| 25 |
+
CriticSubagent,
|
| 26 |
+
PlannerSubagent,
|
| 27 |
+
WorldModelerSubagent,
|
| 28 |
+
)
|
| 29 |
+
from tests._helpers.llm_stub import StubLLMClient
|
| 30 |
+
|
| 31 |
+
# ============================================================================
|
| 32 |
+
# Test fixtures
|
| 33 |
+
# ============================================================================
|
| 34 |
+
|
| 35 |
+
|
| 36 |
+
_VALID_BELIEF_PAYLOAD: Dict[str, Any] = {
|
| 37 |
+
"brain": "epidemiology",
|
| 38 |
+
"latent_estimates": {
|
| 39 |
+
"R1": {
|
| 40 |
+
"estimated_infection_rate": 0.05,
|
| 41 |
+
"estimated_r_effective": 1.2,
|
| 42 |
+
"estimated_compliance": 0.85,
|
| 43 |
+
"confidence_intervals": {},
|
| 44 |
+
},
|
| 45 |
+
},
|
| 46 |
+
"hypotheses": [{"label": "rising-r1", "weight": 0.6, "explanation": "telemetry trending up"}],
|
| 47 |
+
"uncertainty": 0.4,
|
| 48 |
+
"reducible_by_more_thought": 0.3,
|
| 49 |
+
"evidence": [
|
| 50 |
+
{"source": "telemetry", "ref": "R1.reported_cases@t3", "excerpt": "rising"},
|
| 51 |
+
{"source": "policy", "ref": "R1.restriction", "excerpt": "moderate"},
|
| 52 |
+
],
|
| 53 |
+
}
|
| 54 |
+
|
| 55 |
+
|
| 56 |
+
_VALID_PLAN_PAYLOAD: Dict[str, Any] = {
|
| 57 |
+
"action_sketch": "Deploy 100 test_kits to R1",
|
| 58 |
+
"expected_outer_action": {
|
| 59 |
+
"kind": "deploy_resource",
|
| 60 |
+
"region": "R1",
|
| 61 |
+
"resource_type": "test_kits",
|
| 62 |
+
"quantity": 100,
|
| 63 |
+
},
|
| 64 |
+
"expected_value": 0.6,
|
| 65 |
+
"cost": 200.0,
|
| 66 |
+
"assumptions": ["kits inventory > 100"],
|
| 67 |
+
"falsifiers": ["R1 cases drop without intervention"],
|
| 68 |
+
"confidence": 0.75,
|
| 69 |
+
}
|
| 70 |
+
|
| 71 |
+
|
| 72 |
+
_VALID_CRITIC_PAYLOAD: Dict[str, Any] = {
|
| 73 |
+
"brain": "logistics",
|
| 74 |
+
"target_plan_id": "plan-1",
|
| 75 |
+
"attacks": ["ignores R3 hospital saturation"],
|
| 76 |
+
"missing_considerations": ["compliance decay over 4 ticks"],
|
| 77 |
+
"would_change_mind_if": ["new R3 telemetry shows under-utilisation"],
|
| 78 |
+
"severity": 0.6,
|
| 79 |
+
}
|
| 80 |
+
|
| 81 |
+
|
| 82 |
+
def _valid_json_for(role: str, brain: str = "epidemiology") -> str:
|
| 83 |
+
"""Return a valid JSON-schema response for ``role``, brain-substituted."""
|
| 84 |
+
if role == "world_modeler":
|
| 85 |
+
payload = dict(_VALID_BELIEF_PAYLOAD)
|
| 86 |
+
payload["brain"] = brain
|
| 87 |
+
return json.dumps(payload)
|
| 88 |
+
if role == "planner":
|
| 89 |
+
return json.dumps(_VALID_PLAN_PAYLOAD)
|
| 90 |
+
if role == "critic":
|
| 91 |
+
payload = dict(_VALID_CRITIC_PAYLOAD)
|
| 92 |
+
payload["brain"] = brain
|
| 93 |
+
return json.dumps(payload)
|
| 94 |
+
raise ValueError(f"unknown role: {role}")
|
| 95 |
+
|
| 96 |
+
|
| 97 |
+
def _make_subagent_input(
|
| 98 |
+
brain: str = "epidemiology",
|
| 99 |
+
role: str = "world_modeler",
|
| 100 |
+
tick: int = 3,
|
| 101 |
+
round_: int = 1,
|
| 102 |
+
target_plan_id: Optional[str] = None,
|
| 103 |
+
) -> SubagentInput:
|
| 104 |
+
"""Minimal valid SubagentInput. ``round_`` arg avoids shadowing builtin."""
|
| 105 |
+
return SubagentInput(
|
| 106 |
+
brain=brain,
|
| 107 |
+
role=role,
|
| 108 |
+
tick=tick,
|
| 109 |
+
round=round_,
|
| 110 |
+
perception=PerceptionReport(
|
| 111 |
+
brain=brain,
|
| 112 |
+
salient_signals=["R1 cases rising"],
|
| 113 |
+
anomalies=[],
|
| 114 |
+
confidence=0.7,
|
| 115 |
+
evidence=[EvidenceCitation(source="telemetry", ref="R1.cases", excerpt="rising")],
|
| 116 |
+
),
|
| 117 |
+
prior_belief=None,
|
| 118 |
+
prior_plans=[],
|
| 119 |
+
target_plan_id=target_plan_id,
|
| 120 |
+
last_reward=0.5,
|
| 121 |
+
recent_action_log_excerpt=[],
|
| 122 |
+
)
|
| 123 |
+
|
| 124 |
+
|
| 125 |
+
# ============================================================================
|
| 126 |
+
# T1 - WorldModeler emits BeliefState
|
| 127 |
+
# ============================================================================
|
| 128 |
+
|
| 129 |
+
|
| 130 |
+
def test_world_modeler_emits_belief_state() -> None:
|
| 131 |
+
stub = StubLLMClient(scripted_responses=[_valid_json_for("world_modeler", "epidemiology")])
|
| 132 |
+
agent = WorldModelerSubagent(llm_client=stub)
|
| 133 |
+
input_data = _make_subagent_input(brain="epidemiology", role="world_modeler")
|
| 134 |
+
|
| 135 |
+
result = agent.run(input_data, step_idx=0)
|
| 136 |
+
|
| 137 |
+
assert isinstance(result, BeliefState)
|
| 138 |
+
assert result.brain == "epidemiology"
|
| 139 |
+
assert "R1" in result.latent_estimates
|
| 140 |
+
assert len(result.evidence) >= 1
|
| 141 |
+
assert stub.call_count == 1
|
| 142 |
+
|
| 143 |
+
|
| 144 |
+
# ============================================================================
|
| 145 |
+
# T2 - Planner emits CandidatePlan
|
| 146 |
+
# ============================================================================
|
| 147 |
+
|
| 148 |
+
|
| 149 |
+
def test_planner_emits_candidate_plan() -> None:
|
| 150 |
+
stub = StubLLMClient(scripted_responses=[_valid_json_for("planner", "epidemiology")])
|
| 151 |
+
agent = PlannerSubagent(llm_client=stub)
|
| 152 |
+
input_data = _make_subagent_input(brain="epidemiology", role="planner")
|
| 153 |
+
|
| 154 |
+
result = agent.run(input_data, step_idx=1)
|
| 155 |
+
|
| 156 |
+
assert isinstance(result, CandidatePlan)
|
| 157 |
+
assert result.expected_outer_action.kind == "deploy_resource"
|
| 158 |
+
assert result.confidence == 0.75
|
| 159 |
+
|
| 160 |
+
|
| 161 |
+
# ============================================================================
|
| 162 |
+
# T3 - Critic emits CriticReport
|
| 163 |
+
# ============================================================================
|
| 164 |
+
|
| 165 |
+
|
| 166 |
+
def test_critic_emits_critic_report() -> None:
|
| 167 |
+
stub = StubLLMClient(scripted_responses=[_valid_json_for("critic", "logistics")])
|
| 168 |
+
agent = CriticSubagent(llm_client=stub)
|
| 169 |
+
input_data = _make_subagent_input(brain="logistics", role="critic", target_plan_id="plan-1")
|
| 170 |
+
|
| 171 |
+
result = agent.run(input_data, step_idx=2)
|
| 172 |
+
|
| 173 |
+
assert isinstance(result, CriticReport)
|
| 174 |
+
assert result.brain == "logistics"
|
| 175 |
+
assert result.target_plan_id == "plan-1"
|
| 176 |
+
assert result.severity == 0.6
|
| 177 |
+
|
| 178 |
+
|
| 179 |
+
# ============================================================================
|
| 180 |
+
# T4 - Parse failure then retry succeeds (2 LLM calls)
|
| 181 |
+
# ============================================================================
|
| 182 |
+
|
| 183 |
+
|
| 184 |
+
def test_subagent_parse_failure_then_retry_succeeds() -> None:
|
| 185 |
+
stub = StubLLMClient(
|
| 186 |
+
scripted_responses=["not-json-garbage", _valid_json_for("world_modeler", "epidemiology")]
|
| 187 |
+
)
|
| 188 |
+
agent = WorldModelerSubagent(llm_client=stub)
|
| 189 |
+
|
| 190 |
+
result = agent.run(_make_subagent_input(), step_idx=0)
|
| 191 |
+
|
| 192 |
+
assert isinstance(result, BeliefState)
|
| 193 |
+
assert result.brain == "epidemiology"
|
| 194 |
+
assert stub.call_count == 2, "expected 1 initial call + 1 retry"
|
| 195 |
+
|
| 196 |
+
|
| 197 |
+
# ============================================================================
|
| 198 |
+
# T5 - Parse failure twice -> empty fallback (no third LLM call)
|
| 199 |
+
# ============================================================================
|
| 200 |
+
|
| 201 |
+
|
| 202 |
+
def test_subagent_parse_failure_then_retry_fails_returns_empty() -> None:
|
| 203 |
+
stub = StubLLMClient(scripted_responses=["garbage-1", "garbage-2"])
|
| 204 |
+
agent = WorldModelerSubagent(llm_client=stub)
|
| 205 |
+
|
| 206 |
+
result = agent.run(_make_subagent_input(), step_idx=0)
|
| 207 |
+
|
| 208 |
+
assert isinstance(result, BeliefState)
|
| 209 |
+
assert result.brain == "epidemiology"
|
| 210 |
+
assert result.latent_estimates == {}
|
| 211 |
+
assert result.hypotheses == []
|
| 212 |
+
assert result.evidence == []
|
| 213 |
+
assert result.uncertainty == 1.0
|
| 214 |
+
assert result.reducible_by_more_thought == 0.0
|
| 215 |
+
assert stub.call_count == 2, "must NOT make a 3rd call after retry failure"
|
| 216 |
+
|
| 217 |
+
|
| 218 |
+
# ============================================================================
|
| 219 |
+
# T6 - caller_id format matches Phase A Decision 7
|
| 220 |
+
# ============================================================================
|
| 221 |
+
|
| 222 |
+
|
| 223 |
+
_CALLER_ID_RE = re.compile(
|
| 224 |
+
r"^cortex:(epidemiology|logistics|governance):"
|
| 225 |
+
r"(world_modeler|planner|critic):"
|
| 226 |
+
r"t\d+:r[12]:s\d+$"
|
| 227 |
+
)
|
| 228 |
+
|
| 229 |
+
|
| 230 |
+
@pytest.mark.parametrize(
|
| 231 |
+
"role_cls,brain,role_name",
|
| 232 |
+
[
|
| 233 |
+
(WorldModelerSubagent, "epidemiology", "world_modeler"),
|
| 234 |
+
(PlannerSubagent, "logistics", "planner"),
|
| 235 |
+
(CriticSubagent, "governance", "critic"),
|
| 236 |
+
],
|
| 237 |
+
)
|
| 238 |
+
def test_subagent_caller_id_format(
|
| 239 |
+
role_cls: type,
|
| 240 |
+
brain: str,
|
| 241 |
+
role_name: str,
|
| 242 |
+
) -> None:
|
| 243 |
+
stub = StubLLMClient(scripted_responses=[_valid_json_for(role_name, brain)])
|
| 244 |
+
agent = role_cls(llm_client=stub)
|
| 245 |
+
input_data = _make_subagent_input(
|
| 246 |
+
brain=brain,
|
| 247 |
+
role=role_name,
|
| 248 |
+
tick=7,
|
| 249 |
+
round_=2,
|
| 250 |
+
target_plan_id="plan-X" if role_name == "critic" else None,
|
| 251 |
+
)
|
| 252 |
+
|
| 253 |
+
agent.run(input_data, step_idx=4)
|
| 254 |
+
|
| 255 |
+
caller_id = stub.calls[0].caller_id
|
| 256 |
+
assert _CALLER_ID_RE.match(caller_id), (
|
| 257 |
+
f"caller_id={caller_id!r} does not match the locked format"
|
| 258 |
+
)
|
| 259 |
+
assert caller_id == f"cortex:{brain}:{role_name}:t7:r2:s4"
|
| 260 |
+
|
| 261 |
+
|
| 262 |
+
# ============================================================================
|
| 263 |
+
# T8 - SYS prompt loaded from prompts/<role>.txt and brain-formatted
|
| 264 |
+
# (folds in the prompt-formatting refinement: format() must not raise)
|
| 265 |
+
# ============================================================================
|
| 266 |
+
|
| 267 |
+
|
| 268 |
+
@pytest.mark.parametrize(
|
| 269 |
+
"role_cls,role_name",
|
| 270 |
+
[
|
| 271 |
+
(WorldModelerSubagent, "world_modeler"),
|
| 272 |
+
(PlannerSubagent, "planner"),
|
| 273 |
+
(CriticSubagent, "critic"),
|
| 274 |
+
],
|
| 275 |
+
)
|
| 276 |
+
def test_subagent_uses_loaded_prompt_from_file(role_cls: type, role_name: str) -> None:
|
| 277 |
+
raw = (PROMPTS_DIR / f"{role_name}.txt").read_text(encoding="utf-8")
|
| 278 |
+
|
| 279 |
+
# Refinement: format() must not raise even with extra kwargs ignored.
|
| 280 |
+
# Catches {{/}}-escape regressions in JSON-schema sections of the prompts.
|
| 281 |
+
formatted = raw.format(brain="epidemiology", target_plan_id="plan-X")
|
| 282 |
+
assert isinstance(formatted, str)
|
| 283 |
+
assert "epidemiology" in formatted
|
| 284 |
+
|
| 285 |
+
stub = StubLLMClient(scripted_responses=[_valid_json_for(role_name, "epidemiology")])
|
| 286 |
+
agent = role_cls(llm_client=stub)
|
| 287 |
+
input_data = _make_subagent_input(
|
| 288 |
+
brain="epidemiology",
|
| 289 |
+
role=role_name,
|
| 290 |
+
target_plan_id="plan-X" if role_name == "critic" else None,
|
| 291 |
+
)
|
| 292 |
+
|
| 293 |
+
agent.run(input_data, step_idx=0)
|
| 294 |
+
|
| 295 |
+
sys_msg = stub.calls[0].messages[0]
|
| 296 |
+
assert sys_msg.role == "system"
|
| 297 |
+
assert sys_msg.content == formatted
|
| 298 |
+
|
| 299 |
+
|
| 300 |
+
# ============================================================================
|
| 301 |
+
# T9 - Token counter is billed to the expected caller_id
|
| 302 |
+
# ============================================================================
|
| 303 |
+
|
| 304 |
+
|
| 305 |
+
def test_subagent_token_counter_billed_correctly() -> None:
|
| 306 |
+
stub = StubLLMClient(scripted_responses=[_valid_json_for("world_modeler", "epidemiology")])
|
| 307 |
+
agent = WorldModelerSubagent(llm_client=stub)
|
| 308 |
+
input_data = _make_subagent_input(brain="epidemiology", role="world_modeler", tick=3, round_=1)
|
| 309 |
+
|
| 310 |
+
agent.run(input_data, step_idx=0)
|
| 311 |
+
|
| 312 |
+
expected_caller_id = "cortex:epidemiology:world_modeler:t3:r1:s0"
|
| 313 |
+
assert stub.tokens_used_for(expected_caller_id) > 0, (
|
| 314 |
+
"tokens must be billed to the per-role caller_id, not silently lost"
|
| 315 |
+
)
|
| 316 |
+
assert stub.tokens_used_for("never:called") == 0
|
| 317 |
+
|
| 318 |
+
|
| 319 |
+
# ============================================================================
|
| 320 |
+
# T10 - empty_fallback shape locked per Phase A Decisions 6 + 62
|
| 321 |
+
# ============================================================================
|
| 322 |
+
|
| 323 |
+
|
| 324 |
+
def test_subagent_empty_fallback_shape_locked() -> None:
|
| 325 |
+
# WorldModeler: empty BeliefState
|
| 326 |
+
bs = WorldModelerSubagent.empty_fallback("epidemiology")
|
| 327 |
+
assert isinstance(bs, BeliefState)
|
| 328 |
+
assert bs.brain == "epidemiology"
|
| 329 |
+
assert bs.latent_estimates == {}
|
| 330 |
+
assert bs.hypotheses == []
|
| 331 |
+
assert bs.uncertainty == 1.0
|
| 332 |
+
assert bs.reducible_by_more_thought == 0.0
|
| 333 |
+
assert bs.evidence == []
|
| 334 |
+
|
| 335 |
+
# Planner: empty CandidatePlan with NoOp + confidence=0
|
| 336 |
+
cp = PlannerSubagent.empty_fallback("epidemiology")
|
| 337 |
+
assert isinstance(cp, CandidatePlan)
|
| 338 |
+
assert cp.expected_outer_action.kind == "no_op"
|
| 339 |
+
assert cp.expected_value == 0.0
|
| 340 |
+
assert cp.cost == 0.0
|
| 341 |
+
assert cp.assumptions == []
|
| 342 |
+
assert cp.falsifiers == []
|
| 343 |
+
assert cp.confidence == 0.0
|
| 344 |
+
|
| 345 |
+
# Critic: empty CriticReport with severity=0
|
| 346 |
+
cr = CriticSubagent.empty_fallback("epidemiology", target_plan_id="plan-X")
|
| 347 |
+
assert isinstance(cr, CriticReport)
|
| 348 |
+
assert cr.brain == "epidemiology"
|
| 349 |
+
assert cr.target_plan_id == "plan-X"
|
| 350 |
+
assert cr.attacks == []
|
| 351 |
+
assert cr.missing_considerations == []
|
| 352 |
+
assert cr.would_change_mind_if == []
|
| 353 |
+
assert cr.severity == 0.0
|
| 354 |
+
|
| 355 |
+
|
| 356 |
+
# ============================================================================
|
| 357 |
+
# T11 - retry call uses chat-history continuation (sys + usr + bad + retry)
|
| 358 |
+
# ============================================================================
|
| 359 |
+
|
| 360 |
+
|
| 361 |
+
def test_subagent_run_uses_chat_history_on_retry() -> None:
|
| 362 |
+
stub = StubLLMClient(
|
| 363 |
+
scripted_responses=["bad-json", _valid_json_for("world_modeler", "epidemiology")]
|
| 364 |
+
)
|
| 365 |
+
agent = WorldModelerSubagent(llm_client=stub)
|
| 366 |
+
|
| 367 |
+
agent.run(_make_subagent_input(brain="epidemiology", role="world_modeler"), step_idx=0)
|
| 368 |
+
|
| 369 |
+
assert stub.call_count == 2, "expected 2 LLM calls (initial + retry)"
|
| 370 |
+
call1, call2 = stub.calls
|
| 371 |
+
|
| 372 |
+
# call 1 has the original sys + user (2 messages).
|
| 373 |
+
assert len(call1.messages) == 2
|
| 374 |
+
assert call1.messages[0].role == "system"
|
| 375 |
+
assert call1.messages[1].role == "user"
|
| 376 |
+
|
| 377 |
+
# call 2 must contain: original sys + original user + assistant(bad-json) + retry-user.
|
| 378 |
+
# Without chat-history continuation the LLM loses schema context on retry.
|
| 379 |
+
assert len(call2.messages) == 4, "retry must append to the chat history, not start fresh"
|
| 380 |
+
assert call2.messages[0].role == "system"
|
| 381 |
+
assert call2.messages[0].content == call1.messages[0].content
|
| 382 |
+
assert call2.messages[1].role == "user"
|
| 383 |
+
assert call2.messages[1].content == call1.messages[1].content
|
| 384 |
+
assert call2.messages[2].role == "assistant"
|
| 385 |
+
assert call2.messages[2].content == "bad-json"
|
| 386 |
+
assert call2.messages[3].role == "user"
|
| 387 |
+
assert "failed to parse" in call2.messages[3].content.lower()
|
tests/test_outer_reward_in_range.py
CHANGED
|
@@ -1,8 +1,10 @@
|
|
| 1 |
-
"""Outer reward stays in ``[
|
| 2 |
|
| 3 |
-
|
| 4 |
-
|
| 5 |
-
|
|
|
|
|
|
|
| 6 |
|
| 7 |
- Initial state (right after ``load_task``, no ticks applied).
|
| 8 |
- Mid-episode rollouts on all 3 tasks with varied actions.
|
|
@@ -29,7 +31,7 @@ def test_outer_reward_in_range_at_episode_start() -> None:
|
|
| 29 |
for name in TASKS:
|
| 30 |
state = load_task(name, episode_seed=0)
|
| 31 |
r = outer_reward(state, NoOp())
|
| 32 |
-
assert
|
| 33 |
|
| 34 |
|
| 35 |
def test_outer_reward_in_range_during_rollout() -> None:
|
|
@@ -51,8 +53,8 @@ def test_outer_reward_in_range_during_rollout() -> None:
|
|
| 51 |
for action in actions:
|
| 52 |
state = apply_tick(state, action)
|
| 53 |
r = outer_reward(state, action)
|
| 54 |
-
assert
|
| 55 |
-
f"{name} tick={state.tick}: r={r!r} out of [
|
| 56 |
)
|
| 57 |
|
| 58 |
|
|
@@ -67,7 +69,7 @@ def test_outer_reward_in_range_with_high_infection() -> None:
|
|
| 67 |
region.S, region.E, region.I, region.R = 0.0, 0.0, 0.95, 0.05
|
| 68 |
state.tick = state.max_ticks # r_time → 0
|
| 69 |
r = outer_reward(state, NoOp())
|
| 70 |
-
assert
|
| 71 |
|
| 72 |
|
| 73 |
def test_outer_reward_in_range_with_rejected_actions() -> None:
|
|
@@ -81,7 +83,7 @@ def test_outer_reward_in_range_with_rejected_actions() -> None:
|
|
| 81 |
)
|
| 82 |
state = apply_tick(state, a_v2)
|
| 83 |
r = outer_reward(state, a_v2)
|
| 84 |
-
assert
|
| 85 |
|
| 86 |
# Legal-violation: strict severity before escalate-national on hard.
|
| 87 |
state2 = load_task("outbreak_hard", episode_seed=0)
|
|
@@ -89,4 +91,4 @@ def test_outer_reward_in_range_with_rejected_actions() -> None:
|
|
| 89 |
state2 = apply_tick(state2, a_legal)
|
| 90 |
assert state2.recent_action_log[-1].accepted is False
|
| 91 |
r2 = outer_reward(state2, a_legal)
|
| 92 |
-
assert
|
|
|
|
| 1 |
+
"""Outer reward stays in ``[-1.0, 1.0]`` across diverse states.
|
| 2 |
|
| 3 |
+
Range relaxed by Workstream-B Phase-1 fix (M2-A): rejected actions land
|
| 4 |
+
``r_policy = -0.5`` and parse-failure markers land ``r_policy = -1.0``,
|
| 5 |
+
so the per-tick total can go negative. Upper bound stays at 1.0.
|
| 6 |
+
|
| 7 |
+
Covers:
|
| 8 |
|
| 9 |
- Initial state (right after ``load_task``, no ticks applied).
|
| 10 |
- Mid-episode rollouts on all 3 tasks with varied actions.
|
|
|
|
| 31 |
for name in TASKS:
|
| 32 |
state = load_task(name, episode_seed=0)
|
| 33 |
r = outer_reward(state, NoOp())
|
| 34 |
+
assert -1.0 <= r <= 1.0, f"{name}: r={r!r} out of [-1,1] at tick 0"
|
| 35 |
|
| 36 |
|
| 37 |
def test_outer_reward_in_range_during_rollout() -> None:
|
|
|
|
| 53 |
for action in actions:
|
| 54 |
state = apply_tick(state, action)
|
| 55 |
r = outer_reward(state, action)
|
| 56 |
+
assert -1.0 <= r <= 1.0, (
|
| 57 |
+
f"{name} tick={state.tick}: r={r!r} out of [-1,1] after action kind={action.kind!r}"
|
| 58 |
)
|
| 59 |
|
| 60 |
|
|
|
|
| 69 |
region.S, region.E, region.I, region.R = 0.0, 0.0, 0.95, 0.05
|
| 70 |
state.tick = state.max_ticks # r_time → 0
|
| 71 |
r = outer_reward(state, NoOp())
|
| 72 |
+
assert -1.0 <= r <= 1.0, f"high-I worst case: r={r!r}"
|
| 73 |
|
| 74 |
|
| 75 |
def test_outer_reward_in_range_with_rejected_actions() -> None:
|
|
|
|
| 83 |
)
|
| 84 |
state = apply_tick(state, a_v2)
|
| 85 |
r = outer_reward(state, a_v2)
|
| 86 |
+
assert -1.0 <= r <= 1.0, f"V2-rejected: r={r!r}"
|
| 87 |
|
| 88 |
# Legal-violation: strict severity before escalate-national on hard.
|
| 89 |
state2 = load_task("outbreak_hard", episode_seed=0)
|
|
|
|
| 91 |
state2 = apply_tick(state2, a_legal)
|
| 92 |
assert state2.recent_action_log[-1].accepted is False
|
| 93 |
r2 = outer_reward(state2, a_legal)
|
| 94 |
+
assert -1.0 <= r2 <= 1.0, f"legal-violation: r={r2!r}"
|
tests/test_outer_reward_non_constancy.py
CHANGED
|
@@ -15,6 +15,11 @@ from CrisisWorldCortex.models import (
|
|
| 15 |
RestrictMovement,
|
| 16 |
)
|
| 17 |
from CrisisWorldCortex.server.graders import outer_reward
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 18 |
from CrisisWorldCortex.server.simulator import apply_tick, load_task
|
| 19 |
|
| 20 |
|
|
@@ -87,8 +92,12 @@ def test_reward_differs_for_accepted_vs_rejected_action() -> None:
|
|
| 87 |
|
| 88 |
Compare two episodes from the same starting state: one issues NoOp
|
| 89 |
(accepted), the other issues PublicCommunication (V2-rejected). The
|
| 90 |
-
SEIR step runs identically; only ``r_policy`` differs.
|
| 91 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
| 92 |
"""
|
| 93 |
s_a = load_task("outbreak_easy", episode_seed=42)
|
| 94 |
s_b = load_task("outbreak_easy", episode_seed=42)
|
|
@@ -109,7 +118,7 @@ def test_reward_differs_for_accepted_vs_rejected_action() -> None:
|
|
| 109 |
assert r_ok > r_bad, (
|
| 110 |
f"accepted action should score higher than rejected: ok={r_ok!r} bad={r_bad!r}"
|
| 111 |
)
|
| 112 |
-
|
| 113 |
-
assert abs((r_ok - r_bad) -
|
| 114 |
-
f"r_policy gap mismatch: ok-bad={r_ok - r_bad!r}, expected
|
| 115 |
)
|
|
|
|
| 15 |
RestrictMovement,
|
| 16 |
)
|
| 17 |
from CrisisWorldCortex.server.graders import outer_reward
|
| 18 |
+
from CrisisWorldCortex.server.graders.outer_reward import (
|
| 19 |
+
R_POLICY_NOOP_ACCEPTED,
|
| 20 |
+
R_POLICY_REJECTED,
|
| 21 |
+
W_POLICY,
|
| 22 |
+
)
|
| 23 |
from CrisisWorldCortex.server.simulator import apply_tick, load_task
|
| 24 |
|
| 25 |
|
|
|
|
| 92 |
|
| 93 |
Compare two episodes from the same starting state: one issues NoOp
|
| 94 |
(accepted), the other issues PublicCommunication (V2-rejected). The
|
| 95 |
+
SEIR step runs identically; only ``r_policy`` differs. After the
|
| 96 |
+
Workstream-B Phase-1 four-state contract, NoOp(accepted) →
|
| 97 |
+
R_POLICY_NOOP_ACCEPTED (0.0) and PublicCommunication(honesty=0.9,
|
| 98 |
+
rejected as legal-violation) → R_POLICY_REJECTED (-0.5). The exact
|
| 99 |
+
gap is therefore ``(R_POLICY_NOOP_ACCEPTED - R_POLICY_REJECTED) *
|
| 100 |
+
W_POLICY`` (modulo float rounding).
|
| 101 |
"""
|
| 102 |
s_a = load_task("outbreak_easy", episode_seed=42)
|
| 103 |
s_b = load_task("outbreak_easy", episode_seed=42)
|
|
|
|
| 118 |
assert r_ok > r_bad, (
|
| 119 |
f"accepted action should score higher than rejected: ok={r_ok!r} bad={r_bad!r}"
|
| 120 |
)
|
| 121 |
+
expected_gap = (R_POLICY_NOOP_ACCEPTED - R_POLICY_REJECTED) * W_POLICY
|
| 122 |
+
assert abs((r_ok - r_bad) - expected_gap) < 1e-9, (
|
| 123 |
+
f"r_policy gap mismatch: ok-bad={r_ok - r_bad!r}, expected {expected_gap!r}"
|
| 124 |
)
|