# ContextCorruption-Env — Engineering Spec
> OpenEnv Hackathon | Meta × HuggingFace × PyTorch
> Team: Siddh + Teammate | Deadline: Tomorrow 5pm | Budget: $60 HF Credits

---

## Division of Ownership

| Owner | Scope |
|---|---|
| **Siddh** | `environment/` (env, reward, actions, server) + `training/` |
| **Teammate** | `data/` (loader, corruption, generator, facts.json) |
| **Both** | Integration smoke test (Tomorrow 9am), Docker + HF deploy (Tomorrow 11am) |

---

## Repo Layout

```
context-corruption-env/
├── CLAUDE.md
├── README.md
├── openenv.yaml
├── Dockerfile
├── requirements.txt
├── environment/          # Siddh
│   ├── __init__.py
│   ├── actions.py        # Pydantic schemas
│   ├── reward.py         # scoring logic
│   ├── env.py            # OpenEnv Environment subclass
│   └── server.py         # FastAPI app via OpenEnv helper
├── data/                 # Teammate
│   ├── loader.py
│   ├── corruption.py
│   ├── generator.py
│   └── facts.json        # generated artifact, 500+ QA pairs
├── training/
│   ├── train_grpo.py
│   └── ContextCorruption_GRPO.ipynb
├── eval/
│   └── baseline_eval.py
└── assets/
    ├── reward_curve.png
    └── loss_curve.png
```

---

## Bootstrap (both, do this first)

```bash
# Create repo on GitHub: Public, Python .gitignore, MIT license
git clone https://github.com/YOUR_USERNAME/context-corruption-env.git
cd context-corruption-env

mkdir -p environment data training eval assets
touch environment/__init__.py
touch environment/{actions,reward,env,server}.py
touch data/{loader,corruption,generator}.py
touch training/train_grpo.py eval/baseline_eval.py
touch requirements.txt openenv.yaml Dockerfile

git add . && git commit -m "feat: initial structure" && git push origin main

# Siddh's branch
git checkout -b feat/environment && git push origin feat/environment

# Teammate's branch
git checkout -b feat/data-pipeline && git push origin feat/data-pipeline

# Shared venv
python -m venv venv && source venv/bin/activate
pip install openenv fastapi uvicorn websockets pydantic \
            datasets transformers torch trl unsloth \
            wandb faker python-dotenv
pip freeze > requirements.txt
```

---

## environment/actions.py

**Purpose:** Define the shared data contracts. Everything else imports from here.

**ActionType enum** — four string-valued variants:
- `read_doc` — read a document by index (costs budget, no state change)
- `flag_suspicious` — mark a doc as potentially corrupted
- `unflag_doc` — remove a flag
- `submit_answer` — end the episode with a final answer

**ContextCorruptionAction (Pydantic BaseModel):**
- `action_type: ActionType`
- `doc_id: Optional[int]` — 0-indexed, only used for doc actions
- `answer: Optional[str]` — only used on submit
- `confidence: Optional[float]` — 0.0–1.0, only used on submit; validate range

**Document (Pydantic BaseModel):**
- `id: int`, `title: str`, `content: str`
- `is_flagged: bool = False` — this is the *agent's* flag, not ground truth

**EpisodeObservation (Pydantic BaseModel):**
- `question: str`
- `documents: list[Document]`
- `flagged_ids: list[int]`
- `budget_remaining: int`
- `turn: int`
- `episode_done: bool = False`
- `reward: Optional[float]` — only populated after SUBMIT_ANSWER or budget exhaustion
- `message: Optional[str]` — optional human-readable status

---

## environment/reward.py

**Purpose:** Score a completed episode. No LLM judge — fully deterministic.

**Function signature:**
```
compute_reward(submitted_answer, ground_truth_answer, flagged_ids,
               corrupt_ids, confidence, budget_used, max_budget)
    -> tuple[float, dict]
```

**Scoring breakdown (weights sum to ~1.05 max, floor ~-0.5):**

| Component | Logic | Weight |
|---|---|---|
| Answer correctness | normalize both strings (lowercase, strip punct, collapse whitespace), exact match | +0.4 |
| Flag recall | `len(true_positives) / len(corrupt_ids)` — fraction of corrupt docs caught | +0.3 |
| Precision (no false flags) | start at +0.2, subtract 0.1 per false positive, floor at 0 | +0.2 |
| Confidence calibration | if correct: `+0.1 × confidence`; if wrong: `-0.2 × confidence` | ±0.1 |
| Efficiency | `0.05 × (1 - budget_used / max_budget)` — small bonus, don't over-optimise | +0.05 |

Return `(round(total, 4), breakdown_dict)`. The breakdown dict must include all component keys plus `"total"`.

**Private helper:** `_normalize(text: str) -> str` — lowercase, strip non-word chars, collapse whitespace.

---

## environment/env.py

**Purpose:** The stateful RL environment. Subclass `openenv.Environment`.

**Class constants:**
- `MAX_BUDGET = 12`
- `NUM_DOCS = 8`
- `DIFFICULTY_LEVELS = [1, 2, 3, 4]` (number of corrupt docs per episode)

**`__init__(self, difficulty=None)`**
- Store difficulty (None = random per episode)
- Load `data/facts.json` from `Path(__file__).parent.parent / "data" / "facts.json"`
- If the file doesn't exist, fall back to a hardcoded single-item list so the env still imports cleanly
- Call `_reset_state()`

**`reset() -> EpisodeObservation`**
- Call `_reset_state()`
- Sample a random fact
- Pick `n_corrupt`: use `self.difficulty` if set, else `random.choice(DIFFICULTY_LEVELS)`
- Sample `n_corrupt` positions from `range(NUM_DOCS)` without replacement → store as `self._corrupt_ids`
- Call `data.generator.generate_documents(fact, num_docs=NUM_DOCS, corrupt_positions=self._corrupt_ids)`
- Store question, ground truth answer, documents on self
- Return `_build_observation()`

**`step(action: ContextCorruptionAction) -> EpisodeObservation`**
- If `self._done`, return observation with message "Episode already done."
- Increment `self._turn` and `self._budget_used`
- Dispatch on `action.action_type`:
  - `READ_DOC` — no-op (document content is already in the observation; the budget cost is the point)
  - `FLAG_SUSPICIOUS` — append `action.doc_id` to `self._flagged_ids` if not already there
  - `UNFLAG_DOC` — remove `action.doc_id` from `self._flagged_ids` if present
  - `SUBMIT_ANSWER` — call `compute_reward(...)`, set `self._done = True`, store reward + breakdown
- After dispatch, check if `budget_used >= MAX_BUDGET` and not yet done → force submission with empty answer, confidence 0.0
- Return `_build_observation(reward=reward)`

**`state() -> dict`** — return all internal state including ground truth (for logging/debug only, never surfaced to the agent via the observation)

**`_reset_state()`** — zero/clear all instance variables

**`_build_observation(reward=None, message=None) -> EpisodeObservation`** — build the Pydantic observation from current state; set `is_flagged` on each Document based on `self._flagged_ids`

---

## environment/server.py

**Purpose:** Expose the env over HTTP/WebSocket for TRL to connect to.

Use OpenEnv's `create_app` helper — it handles session management and WebSocket routing automatically. You only need to pass:
- `env_factory`: a zero-arg callable that returns a fresh `ContextCorruptionEnv()`
- `action_model`: `ContextCorruptionAction`
- `observation_model`: `EpisodeObservation`
- `max_concurrent_envs`: 64

Assign the result to `app` so uvicorn can find it. Add a `__main__` guard that runs uvicorn on `0.0.0.0:8000`.

> **Note:** If the `create_app` helper name differs in the installed version, check `meta-pytorch/OpenEnv` GitHub for the current API surface.

---

## data/loader.py

**Purpose:** Pull QA facts from three sources, merge, shuffle, write `facts.json`.

**Three loaders (implement as separate functions, called by `build_fact_database`):**

1. **`load_natural_questions(n=300)`**
   - Dataset: `google-research-datasets/natural_questions`, `train` split, streaming=True
   - Filter: only rows where `annotations.short_answers[0].text` exists and has ≤5 words
   - Shape each fact as `{question, answer, source: "natural_questions", conflict_type: "entity"}`

2. **`load_popqa(n=150)`**
   - Dataset: `akariasai/PopQA`, `test` split
   - Filter: rows where `possible_answers` is non-empty
   - Shape: `{question, answer: possible_answers[0], source: "popqa", conflict_type: "entity", entity, relation}`

3. **`load_faitheval_counterfactual(n=100)`**
   - Source: raw JSON from the SalesforceAIResearch/FaithEval GitHub repo (`data/counterfactual.json`)
   - Fetch with `urllib.request`; wrap in try/except and return `[]` on any failure (URL may be down)
   - Shape: `{question, answer, source: "faitheval", conflict_type: "counterfactual", provided_context}`

**`build_fact_database()`** — call all three, concatenate, shuffle, write to `data/facts.json`. Also callable as `__main__`.

---

## data/corruption.py

**Purpose:** Four escalating corruption functions. This is the creative heart of the dataset.

**Level 1 — `corrupt_number(text, answer)`**
- Find all integers in text with regex `\b\d{4}\b|\b\d+\b`
- If the number looks like a year (4 digits, 1900–2030), shift by a random offset (±5, ±10, ±20)
- Otherwise multiply by a random factor (0.5×, 2×, 3×, 5×, 10×)
- Fallback if no numbers found: append a note claiming the figure was revised

**Level 2 — `corrupt_entity(text, answer)`**
- Maintain pools of plausible substitutes by category: countries, cities, person names (use Faker), organizations
- If the answer appears in text, replace it with a different member of the most-fitting pool
- Fallback: append a sentence attributing the fact to a Faker-generated person name

**Level 3 — `corrupt_inversion(text, answer)`**
- Maintain a hardcoded antonym map: largest↔smallest, first↔last, highest↔lowest, won↔lost, north↔south, etc.
- Case-preserving replacement (preserve UPPER/Title/lower casing of the matched word)
- Fallback: append a sentence saying this contradicts earlier scholarly consensus

**Level 4 — `corrupt_coherent(text, answer)`**
- Generate a plausible wrong answer via `_generate_wrong_answer`:
  - If answer contains digits → mutate a number (±1, ±2, ±5)
  - If single capitalised word → use `fake.last_name()`
  - Multi-word → shuffle words
- Insert the wrong answer into the text *and* wrap it in an authoritative-sounding template citing a fake source + fake org + plausible year
- This level should read as convincing to a careful human reader

**Dispatcher:** `corrupt_text(text, answer, level: int) -> str` — route to the right function, catch all exceptions and return a safe fallback.

---

## data/generator.py

**Purpose:** Wrap a raw QA fact into a list of 8 document dicts, some corrupted.

**`generate_documents(fact, num_docs=8, corrupt_positions=None) -> list[dict]`**

For each document index:
- Pick a random source name (e.g. "Encyclopedia Britannica", "Reuters Fact Check", etc. — maintain a pool of ~10)
- Pick a random sentence template that incorporates `{source}`, `{question}`, and `{answer}`
- If the index is in `corrupt_positions`, apply `corrupt_text` at an escalating level (first corrupt doc = level 1, second = level 2, etc., capping at 4)
- Return a dict with `id`, `title`, `content`, `is_corrupt` (ground truth — never shown to the agent)

---

## Integration Smoke Test (Tomorrow 9am — run together)

```python
from environment.env import ContextCorruptionEnv
from environment.actions import ContextCorruptionAction, ActionType

env = ContextCorruptionEnv(difficulty=2)
obs = env.reset()
assert len(obs.documents) == 8
assert obs.budget_remaining == 12

obs = env.step(ContextCorruptionAction(action_type=ActionType.READ_DOC, doc_id=0))
assert obs.budget_remaining == 11

obs = env.step(ContextCorruptionAction(action_type=ActionType.FLAG_SUSPICIOUS, doc_id=0))
assert 0 in obs.flagged_ids

obs = env.step(ContextCorruptionAction(
    action_type=ActionType.SUBMIT_ANSWER, answer="test", confidence=0.8))
assert obs.episode_done
assert -0.5 <= obs.reward <= 1.05
print("Smoke test passed. State:", env.state())
```

This must run without errors before starting training.

---

## training/train_grpo.py

**Purpose:** Fine-tune Qwen2-1.5B-Instruct with GRPO against the live env server.

**Steps to implement:**

1. Init WandB: `wandb.init(project="context-corruption-env", name="qwen-1.5b-grpo-run1")`

2. Load model with Unsloth:
   - `unsloth/Qwen2-1.5B-Instruct`, `max_seq_length=2048`, `load_in_4bit=True`
   - Apply LoRA: `r=16`, target `q/k/v/o_proj`, no dropout, `use_gradient_checkpointing="unsloth"`

3. Configure GRPO:
   - `num_train_epochs=3`
   - `per_device_train_batch_size=4`, `gradient_accumulation_steps=4`
   - `learning_rate=5e-5`, `max_completion_length=512`
   - `num_generations=8` (GRPO group size)
   - `report_to="wandb"`, `logging_steps=10`, `save_steps=50`

4. System prompt for the agent — describe the task, list the three tools (`read_doc`, `flag_suspicious`, `submit_answer`), and give the strategy: cross-reference claims, flag inconsistencies, trust parametric knowledge when docs conflict.

5. `env_factory = lambda: EnvClient(base_url=ENV_URL).sync()` — fill `ENV_URL` after HF Space is live.

6. Construct `GRPOTrainer(model, config, environment_factory, tokenizer)` and call `.train()`.

7. After training: `wandb.finish()`. Download reward + loss plots from WandB and save as `assets/reward_curve.png` and `assets/loss_curve.png`. Commit them.

**Budget guidance:** ~200 steps target; minimum 50 steps with visible upward trend is acceptable. Reserve ~$20 of credits for reruns.

---

## eval/baseline_eval.py

**Purpose:** Establish a pre-training baseline to make the improvement curves meaningful.

Run 100 episodes with a random agent (randomly flags 0–4 docs, submits "unknown" with 0.5 confidence). Print avg/min/max reward and write to `eval/baseline_results.json`. Run this *before* training.

---

## Dockerfile

```dockerfile
FROM python:3.11-slim
WORKDIR /app
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt
COPY . .
RUN python -c "from data.loader import build_fact_database; build_fact_database()" || true
EXPOSE 7860
CMD ["uvicorn", "environment.server:app", "--host", "0.0.0.0", "--port", "7860"]
```

---

## openenv.yaml

```yaml
name: context-corruption-env
version: "1.0.0"
description: >
  OpenEnv environment for training epistemic robustness in LLMs.
  Agents identify correct answers and flag corrupted documents
  in a multi-doc QA setting with verifiable, objective rewards.
author: "Your Team Name"
license: MIT

environment:
  entrypoint: "environment.server:app"
  action_schema: "environment.actions.ContextCorruptionAction"
  observation_schema: "environment.actions.EpisodeObservation"
  max_concurrent_sessions: 64

reward:
  type: "objective"
  range: [-0.5, 1.05]
  components:
    - {name: answer_correctness,      weight: 0.4}
    - {name: corruption_detection,    weight: 0.3}
    - {name: false_positive_penalty,  weight: 0.2}
    - {name: confidence_calibration,  weight: 0.1}

datasets:
  - {name: Natural Questions, url: "https://huggingface.co/datasets/google-research-datasets/natural_questions"}
  - {name: PopQA,             url: "https://huggingface.co/datasets/akariasai/PopQA"}
  - {name: FaithEval,         url: "https://github.com/SalesforceAIResearch/FaithEval"}

citation: |
  @misc{contextcorruption2026,
    title={ContextCorruption-Env: Training Epistemic Robustness in LLMs},
    year={2026},
    note={OpenEnv Hackathon Submission}
  }
```

---

## HuggingFace Deployment

```bash
pip install huggingface_hub
huggingface-cli login
huggingface-cli repo create context-corruption-env --type space --space_sdk docker
git remote add space https://huggingface.co/spaces/YOUR_HF_USERNAME/context-corruption-env
git push space main
```

Emergency fallback if Space deploy fails: `pip install pyngrok && ngrok http 8000`

---

## Sync Schedule

| Time | Checkpoint |
|------|------------|
| Tonight 12am | Siddh: `env.py` skeleton runs. Teammate: `facts.json` with 100+ entries. |
| Tonight 3am | Full integration attempt — `env.reset()` works with real data. |
| Tomorrow 9am | Smoke test passes. Both push to main. |
| Tomorrow 11am | HF Space live. Training starts. |
| Tomorrow 2pm | Training done. Plots committed. README filled in. |
| Tomorrow 4pm | Final submission check. All links in README. |

---

## Emergency Fallbacks

- **FaithEval URL down:** Skip it — NQ + PopQA gives 450+ facts, enough.
- **Training too slow:** 50 steps with visible upward trend beats no training evidence.
- **OpenEnv API changed:** Check `meta-pytorch/OpenEnv` GitHub for the current `create_app` / `Environment` import paths.

---

## README Template (fill in after training)

Key sections judges want to see:
1. **The Problem** — LLMs defer to wrong retrieved docs even when they know the answer. Standard RLHF makes this worse (cite ClashEval NeurIPS 2024, CANOE May 2025, Knowledgeable-R1 June 2025).
2. **What We Built** — OpenEnv RL environment, objective reward, no LLM judge.
3. **Results table** — Random baseline vs trained model: avg reward, answer accuracy, corruption detection rate.
4. Reward + loss curve images.
5. Links: HF Space, Colab, blog post, WandB run.

---

*OpenEnv Hackathon, April 2026*