docs/agent.md

# TRACE v1 Spec — OpenEnv Incident Response Environment

**Status:** Build-ready production v1  
**Owner:** Rajarshi Datta  
**Timeline:** 7 days  
**Target:** Meta × PyTorch × Hugging Face OpenEnv Hackathon

This spec incorporates critical feedback on the v2.0 PRD. It is **narrowed, execution-ready, and removes all ambiguities.**

---

## Executive Summary

TRACE is a **deterministic, partial-observability RL environment** for **incident response in production infrastructure**. An AI agent interacts with realistic infrastructure incidents by observing systems, running diagnostic actions, and executing remediation. The environment is:

- **OpenEnv-compliant** (pyproject.toml, server/app.py, openenv.yaml)
- **Deterministic** (3 hand-crafted scenarios)
- **Partially observable** (ground truth hidden behind `inspect_*` actions)
- **Action-structured** (action_type + target + value)

- **Outcome-graded** (no diagnosis_accuracy; only resolution success + efficiency)

**Verdict:** This v1 is buildable, complies with validator, and remains challenging.

---

## 1. Problem Statement

Production engineers spend significant time on:

1. **Triage** — filtering false positives from real alerts
2. **Inspection** — digging through logs and metrics
3. **Diagnosis** — identifying root cause
4. **Remediation** — executing fixes (scale, restart, rollback)
5. **Validation** — confirming recovery

Current RL benchmarks do **not** simulate this workflow. TRACE fills that gap.

---

## 2. Design Principles (v1)

### P1 — Partial Observability (FIX #1)

**Previous problem:** Observations exposed `db_status`, `worker_health`, `recent_logs`, `alerts` directly. This leaked too much ground truth.

**Fix:** Observation shows only:
- Generic telemetry (CPU, memory, latency, error_rate, queue_depth)
- Alert names (no context)
- Service status enums (healthy, degraded, down)

Ground truth details (logs, detailed metrics, alert context) are hidden behind inspection actions.

### P2 — Deterministic Scenarios

Exactly 3 hand-crafted incident types, all **reproducible**:

| Scenario        | Root Cause              | Typical Fix              |
|-----------------|-------------------------|--------------------------|
| easy_cpu_spike  | Worker overload         | scale_workers            |

| medium_cascade  | Queue deadlock cascades | restart_service          |

| hard_mixed      | DB + release regression | restart_database + wait  |



### P3 — Action Structure (FIX #2)



**Previous problem:** Actions had no target or magnitude (`restart_service`, `scale_workers` with no arity).



**Fix:** All actions use **triple format:**



```python

(action_type, target, value)
```



Examples:

- `("restart_service", "api_workers", None)`

- `("scale_workers", "api_workers", 5)`

- `("inspect_logs", "database", None)`



### P4 — Reward: Cumulative + Normalized (FIX #3)



**Previous problem:** Rewards clamped to [0,1] per step, causing penalties to collapse to 0.



**Fix:** 

- Collect all step rewards (no per-step clamping)

- Normalize **only at episode end**

- Ensures agent learns long-horizon causality



### P5 — Discovery Action for Diagnosis (FIX #4)



**Previous problem:** Grader includes `diagnosis_accuracy`, but action space has no way to state a diagnosis.



**Fix:** Remove `diagnosis_accuracy` from final grade. Grade only:

- **Resolution success** (binary: incident resolved or not)

- **Efficiency** (steps vs max_steps)



Agent learns diagnosis implicitly through remediation actions.



---



## 3. Environment Architecture



```
┌──────────────────────────────┐
│     inference.py             │
│    (LLM Agent Loop)          │
└──────────┬───────────────────┘
           │ HTTP

    ┌──────┴──────┐

    ▼             ▼

POST /step   GET /state

POST /reset   GET /health

    │             ▲

    └──────┬──────┘

           ▼

    ┌─────────────┐

    │  TraceEnv   │

    │ (gym-like)  │

    └──────┬──────┘

           │

    ┌──────┴──────┬──────────┬─────────┐

    ▼             ▼          ▼         ▼

 scenarios   simulator    rewards   graders

```


---

## 4. Observation Space

```python

class Observation(BaseModel):

    timestamp: str                      # ISO8601

    

    # Metrics (always visible)

    cpu_usage_pct: float               # [0, 100]

    memory_usage_pct: float

    error_rate_pct: float

    api_latency_ms: float

    queue_depth: int

    

    # Service status (always visible, generic)

    services: dict[str, str]           # e.g., {"api_workers": "healthy"}

    

    # Alerts (names only, no context)

    active_alerts: list[str]           # e.g., ["alert_001", "alert_002"]

    

    # Inspection results (populated by inspect_* actions)

    last_inspection: Optional[dict]    # {"type": "logs", "target": "api_workers", "data": "..."}

```

**Key:** Root cause is hidden until agent calls `inspect_logs`, `inspect_metrics`, `inspect_alert`.

---

## 5. Action Space

```python

class Action(BaseModel):

    action_type: str

    target: Optional[str]   # service/metric/alert_id

    value: Optional[float]  # scaling factor, count, etc.

```

**Valid actions:**

| Action | Target | Value | Effect |
|--------|--------|-------|--------|
| `inspect_logs` | service_name | None | Returns log snippet (reveals cause) |

| `inspect_metrics` | metric_name | None | Returns metric timeseries |

| `inspect_alert` | alert_id | None | Returns alert details |

| `restart_service` | service_name | None | Resets service state |

| `scale_workers` | service_name | worker_count | Scales horizontally |
| `restart_database` | None | None | Resets DB state |
| `rollback_release` | None | None | Undoes recent deployment |
| `clear_queue` | None | None | Clears backlog |
| `declare_healthy` | None | None | Declare incident resolved (terminal) |
| `declare_unfixable` | None | None | Give up (terminal) |

---

## 6. Scenario Design

### Scenario 1: `easy_cpu_spike`

**Difficulty:** Beginner (2–4 steps)

**Trigger:** Sudden traffic spike floods API workers.

**Observable symptoms:**
- `cpu_usage_pct` → 85%
- `api_latency_ms` → 500ms
- `error_rate_pct` → 5%
- `active_alerts` → ["alert_cpu_high"]
- `services.api_workers` → "degraded"

**Hidden root cause:** Workload surge, solvable by horizontal scaling

**Optimal trajectory:**
```

1. Observe metrics (CPU high is visible)

2. inspect_logs("api_workers") → reveals "traffic spike, need more workers"

3. scale_workers("api_workers", 5) → CPU → 60%, incident recovers

4. declare_healthy() → DONE

```

**Reward:** Inspection (+1), Remediation (+5), Declare (+10) = success

---

### Scenario 2: `medium_cascade`



**Difficulty:** Intermediate (3–6 steps)



**Trigger:** Queue service memory leak + cascading worker failures.



**Observable symptoms (evolve over steps):**

- Step 1: `queue_depth` rising slowly
- Step 3: `queue_depth` > 500, `error_rate_pct` rising
- Step 5: `services.queue_service` → "degraded", worker timeouts begin
- Step 7: Multiple services → "degraded"

**Hidden root cause:** Queue memory leak; fixable by restart

**Optimal trajectory:**
```

1. Observe metrics (queue_depth unusual)

2. inspect_metrics("queue_depth") → "backlog critical"

3. inspect_logs("queue_service") → "memory usage high, leak suspected"

4. restart_service("queue_service") → queue resets, backlog clears

5. declare_healthy()

```

**Reward:** 2× Inspection (+2), Remediation (+5), Declare (+10) = strong success

---

### Scenario 3: `hard_mixed`



**Difficulty:** Advanced (4–8 steps)



**Trigger:** Recent release + DB connection pool exhaustion + cascading errors.



**Observable symptoms:**

- `error_rate_pct` spiking (5% → 20%)

- `api_latency_ms` very high (100 → 2000ms)

- Multiple alerts: `["alert_high_error_rate", "alert_db_slow", "alert_pool_exhaustion"]`
- `services.database` → "degraded"
- False lead: CPU is high (symptom, not cause)

**Hidden root cause:** DB pool exhausted (release added inefficient queries + not enough connections)

**Optimal trajectory:**
```

1. Observe metrics (error spike, latency spike)

2. inspect_alert("alert_pool_exhaustion") → "DB connection pool at 100%"

3. inspect_logs("database") → "recent release queries inefficient"

4. inspect_metrics("db_connections") → confirms pool exhaustion

5. restart_database() → pool resets, errors drop

6. [optional] rollback_release() if still degraded → teaches causality

7. declare_healthy()

```

**Reward:** 3+ Inspections (+3), Remediation (+8), Declare (+10) = strong success

---

## 7. Reward Structure (FIXED)

### Step-wise Rewards (Accumulated, No Per-Step Clamping)

```python

reward = 0



# Inspection

if action == inspect_logs and target is relevant:

    reward += 1.0

if action == inspect_metrics and target is relevant:

    reward += 1.0

if action == inspect_alert:

    reward += 0.5



# Remediation

if action solves active problem:

    reward += 5.0



# Penalties

if action is duplicate_recent:

    reward -= 0.5

if action worsens incident:

    reward -= 2.0

if action is irrelevant:

    reward -= 0.1



# Terminal

if declare_healthy() and incident_resolved:

    reward += 10.0

if declare_healthy() and NOT incident_resolved:

    reward -= 5.0

```

**All rewards summed across episode. No clamping until end.**

### Final Score (Outcome-Based)

```python

# Normalize accumulated reward

episode_reward = sum(step_rewards) / max_possible_reward

final_reward = min(max(episode_reward, 0), 1.0)



# Grading (NO diagnosis_accuracy)

score = (

    0.6 * (1.0 if incident_resolved else 0.0)  # binary success

    + 0.4 * (1.0 - steps_taken / max_steps)     # efficiency

)

```

**Example:** 
- Easy task: max_steps=5, agent solves in 3 → score = 0.6×1.0 + 0.4×(1 - 3/5) = 0.76

- Hard task: max_steps=8, agent solves in 8 → score = 0.6×1.0 + 0.4×(1 - 8/8) = 0.60

---

## 8. State Transition Logic

### Deterministic Stepping

Each episode uses a **scenario_clock** that progresses deterministically. Same seed → same trajectory.



```python

def transition(state, action) -> (next_state, reward, done):

    # Advance time

    state.timestamp = increment_time(state.seed)

    

    # Apply scenario progression (if no action taken)

    if action not relevant:

        state = apply_scenario_step(state)  # e.g., queue_depth grows

    

    # If remediation action, apply fix

    if action == restart_service:

        state.services[target] = "healthy"

        state = reset_related_metrics()

    

    # Check terminal

    if incident_resolved_enough():

        done = True

    

    return state, reward, done

```



**Property:** `transition(state, a, seed=42)` is deterministic.



---



## 9. API Routes



### POST /reset



Request:

```json

{

    "task_id": "easy_cpu_spike" | "medium_cascade" | "hard_mixed",

    "seed": 42

}

```



Response:

```json

{

    "observation": {...},

    "info": {

        "task_id": "easy_cpu_spike",

        "episode_id": "uuid",

        "max_steps": 5,

        "root_cause": "traffic_spike"

    }

}

```



### POST /step



Request:

```json

{

    "action": {

        "action_type": "scale_workers",

        "target": "api_workers",

        "value": 5

    }

}

```



Response:

```json

{

    "observation": {...},

    "reward": 5.0,

    "done": false,

    "info": {

        "step": 1,

        "episode_reward": 5.0,

        "message": "Workers scaled to 5"

    }

}

```



### GET /state



Response:

```json

{

    "observation": {...},

    "episode_reward": 5.0,

    "steps": 1,

    "done": false

}

```



### GET /health



Response:

```json

{

    "status": "healthy",

    "version": "0.1.0"

}

```



---



## 10. Project Structure



```text

TRACE/

├── pyproject.toml              # OpenEnv spec

├── uv.lock                     # Dependencies locked

├── README.md

├── openenv.yaml                # Environment metadata

├── Dockerfile

├── requirements.txt

├── .env.example

│

├── trace/                      # Core module

│   ├── __init__.py

│   ├── env.py                  # TraceEnv class

│   ├── models.py               # Pydantic schemas

│   ├── scenarios.py            # Scenario generators

│   ├── simulator.py            # State transitions

│   ├── rewards.py              # Reward engine

│   ├── graders.py              # Grading logic

│   └── utils.py                # Helpers

│

├── server/                     # FastAPI app

│   ├── __init__.py

│   └── app.py                  # Routes + server

│

├── inference.py                # Agent policy loop

│

├── tests/

│   ├── __init__.py

│   ├── test_env.py

│   ├── test_api.py

│   ├── test_rewards.py

│   ├── test_graders.py

│   └── test_scenarios.py

│

└── scripts/

    └── run_benchmark.py        # Local evaluation

```



---



## 11. Testing



### Unit Tests



1. **test_scenarios.py:** Seed determinism — verify same seed produces same trajectory
2. **test_rewards.py:** Cumulative rewards (no per-step clamping)

3. **test_graders.py:** Final score calculation
4. **test_env.py:** State transitions



### API Tests



1. POST /reset returns valid Observation

2. POST /step accepts valid Action, returns next state

3. GET /health returns 200 OK

4. Invalid action → 400 Bad Request



### Validation



```bash

openenv validate

./validate-submission.sh

```



---



## 12. Docker & Deployment



```dockerfile

FROM python:3.11-slim

WORKDIR /app

COPY . .

RUN pip install -e .

EXPOSE 7860

CMD ["uvicorn", "server.app:app", "--host", "0.0.0.0", "--port", "7860"]

```



**HF Spaces:** Push to `meta-trace` repo, enable auto-deploy.



---



## 13. Inference Pipeline



**File:** `inference.py`



```python

import os

from openai import OpenAI



client = OpenAI(

    base_url=os.getenv("API_BASE_URL", "http://localhost:7860"),

    api_key=os.getenv("HF_TOKEN")

)



print("[START]")



# Agent loop

response = client.post("/reset", json={"task_id": "easy_cpu_spike", "seed": 0})

state = response.json()["observation"]

done = False



for step in range(MAX_STEPS):

    # LLM decides next action

    action = agent_policy(state)

    

    response = client.post("/step", json={"action": action})

    state = response.json()["observation"]

    reward = response.json()["reward"]

    done = response.json()["done"]

    

    if done:

        break



print("[END]")

```



**Emit exactly:**

- `[START]` before first step

- `[END]` after completion



---



## 14. Risk Register



| Risk | Mitigation |

|------|-----------|

| Validator fails on structure | Continuous `openenv validate` during dev |

| Scenarios become random | Seed-based RNG, determinism tests |

| Reward instability | No per-step clamp, cumulative only |

| Observability too opaque | 3 simple scenarios + dense inspection rewards |

| Diagnosis is ungraded | Removed from final score; implicit in remediation |



---



## 15. Success Criteria (v1 Complete)



✅ `pyproject.toml` + `uv.lock` present  

✅ `openenv validate` passes  

✅ 3 deterministic scenarios reproducible by seed  

✅ API: /reset, /step, /state, /health working  

✅ Rewards cumulative-normalized, no per-step clamp  

✅ Observations hide ground truth (partial observability)  

✅ Actions all use (type, target, value) format  

✅ Grader: 0.6×success + 0.4×efficiency (no diagnosis_accuracy)  

✅ `inference.py` runs, emits `[START]` and `[END]`  

✅ Docker builds and serves  

✅ All tests pass  



---



## 16. Execution Plan (7 Days)



| Day | Milestone |

|-----|-----------|

| 1–2 | Models + scenarios + simulator (determinism verified) |

| 3 | Rewards (cumulative logic) + graders |

| 4 | FastAPI server + Docker + `openenv validate` |

| 5 | `inference.py` + logging + tests |

| 6 | Deploy to HF Spaces |

| 7 | Polish + final validation |



---



**Status:** This spec is **build-ready**. Execute continuously against validator. No further design changes.