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:

(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

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

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)

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)

# 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.

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:

{
    "task_id": "easy_cpu_spike" | "medium_cascade" | "hard_mixed",
    "seed": 42
}

Response:

{
    "observation": {...},
    "info": {
        "task_id": "easy_cpu_spike",
        "episode_id": "uuid",
        "max_steps": 5,
        "root_cause": "traffic_spike"
    }
}

POST /step

Request:

{
    "action": {
        "action_type": "scale_workers",
        "target": "api_workers",
        "value": 5
    }
}

Response:

{
    "observation": {...},
    "reward": 5.0,
    "done": false,
    "info": {
        "step": 1,
        "episode_reward": 5.0,
        "message": "Workers scaled to 5"
    }
}

GET /state

Response:

{
    "observation": {...},
    "episode_reward": 5.0,
    "steps": 1,
    "done": false
}

GET /health

Response:

{
    "status": "healthy",
    "version": "0.1.0"
}

---

10. Project Structure

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

openenv validate
./validate-submission.sh

---

12. Docker & Deployment

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

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.