Upload folder using huggingface_hub

Dockerfile

@@ -0,0 +1,81 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Multi-stage build using openenv-base
+# This Dockerfile is flexible and works for both:
+# - In-repo environments (with local OpenEnv sources)
+# - Standalone environments (with openenv from PyPI/Git)
+# The build script (openenv build) handles context detection and sets appropriate build args.
+
+ARG BASE_IMAGE=ghcr.io/meta-pytorch/openenv-base:latest
+FROM ${BASE_IMAGE} AS builder
+
+WORKDIR /app
+
+# Ensure git is available (required for installing dependencies from VCS)
+RUN apt-get update && \
+    apt-get install -y --no-install-recommends git && \
+    rm -rf /var/lib/apt/lists/*
+
+# Build argument to control whether we're building standalone or in-repo
+ARG BUILD_MODE=in-repo
+ARG ENV_NAME=token_optimiser
+
+# Copy environment code (always at root of build context)
+COPY . /app/env
+
+# For in-repo builds, openenv is already vendored in the build context
+# For standalone builds, openenv will be installed via pyproject.toml
+WORKDIR /app/env
+
+# Ensure uv is available (for local builds where base image lacks it)
+RUN if ! command -v uv >/dev/null 2>&1; then \
+        curl -LsSf https://astral.sh/uv/install.sh | sh && \
+        mv /root/.local/bin/uv /usr/local/bin/uv && \
+        mv /root/.local/bin/uvx /usr/local/bin/uvx; \
+    fi
+    
+# Install dependencies using uv sync
+# If uv.lock exists, use it; otherwise resolve on the fly
+RUN --mount=type=cache,target=/root/.cache/uv \
+    if [ -f uv.lock ]; then \
+        uv sync --frozen --no-install-project --no-editable; \
+    else \
+        uv sync --no-install-project --no-editable; \
+    fi
+
+RUN --mount=type=cache,target=/root/.cache/uv \
+    if [ -f uv.lock ]; then \
+        uv sync --frozen --no-editable; \
+    else \
+        uv sync --no-editable; \
+    fi
+
+# Final runtime stage
+FROM ${BASE_IMAGE}
+
+WORKDIR /app
+
+# Copy the virtual environment from builder
+COPY --from=builder /app/env/.venv /app/.venv
+
+# Copy the environment code
+COPY --from=builder /app/env /app/env
+
+# Set PATH to use the virtual environment
+ENV PATH="/app/.venv/bin:$PATH"
+
+# Set PYTHONPATH so imports work correctly
+ENV PYTHONPATH="/app/env:$PYTHONPATH"
+
+# Health check
+HEALTHCHECK --interval=30s --timeout=3s --start-period=5s --retries=3 \
+    CMD curl -f http://localhost:8000/health || exit 1
+
+# Run the FastAPI server
+# The module path is constructed to work with the /app/env structure
+ENV ENABLE_WEB_INTERFACE=true
+CMD ["sh", "-c", "cd /app/env && uvicorn server.app:app --host 0.0.0.0 --port 8000"]

README.md

@@ -1,10 +1,280 @@
 ---
-title: Token Optimiser
-emoji: 🦀
-colorFrom: red
-colorTo: blue
+title: Token Optimiser Environment
+emoji: 🔤
+colorFrom: indigo
+colorTo: purple
 sdk: docker
+app_port: 8000
 pinned: false
+base_path: /web
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# 🔤 Prompt & Response Token Optimization Environment
+
+> An OpenEnv-compatible RL environment that trains agents to minimize LLM API token usage while preserving semantic quality — reducing AI inference costs at scale.
+
+---
+
+## 📌 Introduction
+
+Large Language Model APIs charge per token. Verbose prompts and unconstrained responses waste tokens and money. This environment trains an AI agent to **rewrite verbose prompts into concise, efficient versions** that:
+
+- Use fewer input tokens
+- Guide the LLM toward shorter, correctly-formatted responses
+- Preserve the full semantic meaning and intent of the original request
+- Respect output format constraints (free text, bullet points, JSON)
+
+The agent learns real-world prompt engineering — a critical skill for production LLM systems where cost efficiency matters at scale.
+
+---
+
+## 🏗️ Architecture
+
+```
+token_optimiser/
+├── inference.py                        # Hackathon evaluation script
+├── models.py                           # Pydantic data models (Action / Observation / State)
+├── client.py                           # Async WebSocket EnvClient
+├── openenv.yaml                        # OpenEnv deployment config
+├── Dockerfile                          # Root-level multi-stage build
+├── pyproject.toml                      # Package config & dependencies
+├── uv.lock                             # Locked dependencies
+└── server/
+    ├── app.py                          # FastAPI server (WebSocket + HTTP)
+    ├── token_optimiser_environment.py  # Core RL environment logic
+    └── requirements.txt                # Server dependencies
+```
+
+---
+
+## 🧠 How It Works
+
+### Action Space
+The agent submits a **`TokenOptimiserAction`** containing its optimized version of the original prompt:
+
+```python
+class TokenOptimiserAction(Action):
+    optimized_prompt: str   # Agent's rewritten, token-efficient prompt
+```
+
+### Observation Space
+After each step, the agent receives a **`TokenOptimiserObservation`**:
+
+```python
+class TokenOptimiserObservation(Observation):
+    llm_response: str    # Actual LLM response to the optimized prompt
+    input_tokens: int    # Token count of the optimized prompt
+    output_tokens: int   # Token count of the LLM response
+    reward: float        # Step reward (0.0 – 1.0)
+```
+
+### State Space
+```python
+class TokenOptimiserState(State):
+    original_prompt: str      # The verbose task prompt the agent must optimize
+    task_difficulty: str      # "easy" | "medium" | "hard"
+    task_index: int           # Index in task bank
+```
+
+---
+
+## 📋 Tasks
+
+### 🟢 Easy — Verbosity Reduction
+**Original:** `"Can you please explain in a very detailed manner what machine learning is and how it works step by step?"`  
+**Goal:** Strip filler words, compress to core query  
+**Expected optimized:** `"Explain machine learning briefly."`  
+**Max output:** 50 tokens
+
+### 🟡 Medium — Format Constraint Addition
+**Original:** `"I need a comprehensive analysis of the renewable energy market trends over the past decade, including solar, wind, and hydroelectric power growth rates..."`  
+**Goal:** Compress input AND add explicit output format constraints  
+**Expected optimized:** `"Summarize 2013-2023 renewable energy trends: solar, wind, hydro. In 5 bullet points."`  
+**Max output:** 100 tokens
+
+### 🔴 Hard — Multi-Intent Structured Output
+**Original:** 82-word complex data science analysis request with 5 sub-tasks  
+**Goal:** Minimize total tokens (input + output) while producing structured JSON with all 5 required keys  
+**Expected optimized:** Compressed prompt specifying `JSON with keys: top_categories, growth_regions, responsive_segments, budget_allocation, risks_watch`  
+**Max output:** 200 tokens
+
+---
+
+## 🏆 Reward Function
+
+**Hybrid grading** — combines token efficiency + LLM-as-judge semantic scoring:
+
+| Component | Weight | Description |
+|-----------|--------|-------------|
+| Token Efficiency | 0.0 – 0.40 | Tokens saved vs. reference (input + output combined) |
+| Semantic Quality | 0.0 – 0.30 | LLM judge rates response quality 0–10 |
+| Format Compliance | 0.0 – 0.20 | Response matches required format (bullets / JSON / brief) |
+| Length Penalty | −0.10 | Output exceeds 2× max token budget |
+
+```
+reward = token_efficiency + (semantic_score × 0.3) + format_score + length_penalty
+reward = clamp(reward, 0.0, 1.0)
+```
+
+| Score | Meaning |
+|-------|---------|
+| 0.0 | Meaning lost, system failure, or no optimization |
+| 0.3 – 0.5 | Token reduction achieved but quality degraded |
+| 0.6 – 0.8 | Good balance of compression and quality |
+| 0.9 – 1.0 | Optimal: minimal tokens, correct format, meaning preserved |
+
+---
+
+## ⚙️ Setup & Installation
+
+### Prerequisites
+- Python 3.10+
+- [`uv`](https://github.com/astral-sh/uv) (recommended) or `pip`
+- A Hugging Face account with a token that has **Inference Providers** permission
+
+### 1. Clone and Install
+
+```bash
+git clone <your-repo-url>
+cd token_optimiser
+
+# Install with uv (recommended)
+uv sync
+
+# Or with pip
+pip install -e .
+```
+
+### 2. Set Environment Variables
+
+```bash
+# Required
+export HF_TOKEN="hf_your_token_here"
+
+# Optional (these are the defaults)
+export API_BASE_URL="https://router.huggingface.co/v1"
+export MODEL_NAME="Qwen/Qwen2.5-72B-Instruct"
+export SERVER_URL="http://localhost:8000"
+```
+
+On Windows (PowerShell):
+```powershell
+$env:HF_TOKEN = "hf_your_token_here"
+```
+
+> **HF Token Permissions:** Your token must have **"Make calls to Inference Providers"** enabled.  
+> Create/edit at → https://huggingface.co/settings/tokens
+
+---
+
+## 🚀 Running Locally
+
+### Step 1 — Start the Environment Server
+
+```bash
+# Terminal 1
+uv run uvicorn server.app:app --host 0.0.0.0 --port 8000
+```
+
+Expected output:
+```
+INFO:     Application startup complete.
+INFO:     Uvicorn running on http://0.0.0.0:8000
+```
+
+### Step 2 — Run the Inference Script
+
+```bash
+# Terminal 2
+uv run inference.py
+```
+
+Expected output:
+```
+[START] task=token_optimization env=token_optimiser model=Qwen/Qwen2.5-72B-Instruct
+[STEP] step=1 action='Analyze renewable energy trends...' reward=0.62 done=false error=null
+[STEP] step=2 action='Summarize 2013-2023 renewable energy...' reward=0.81 done=false error=null
+[STEP] step=3 action='Renewable energy trends 2013-2023...' reward=0.84 done=false error=null
+[STEP] step=4 action='Renewable energy 2013-2023 trends...' reward=0.82 done=false error=null
+[STEP] step=5 action='Energy trends 2013-2023: solar, wind...' reward=0.81 done=true error=null
+[END] success=true steps=5 score=0.780 rewards=0.62,0.81,0.84,0.82,0.81
+```
+
+### Step 3 — (Optional) Run via Docker
+
+```bash
+# Build
+docker build -t token-optimiser-env .
+
+# Run
+docker run -p 8000:8000 \
+  -e HF_TOKEN=$HF_TOKEN \
+  -e API_BASE_URL=$API_BASE_URL \
+  token-optimiser-env
+```
+
+---
+
+## 🐍 Using the Python Client
+
+```python
+import asyncio
+from token_optimiser import TokenOptimiserEnv, TokenOptimiserAction
+
+async def main():
+    async with TokenOptimiserEnv(base_url="http://localhost:8000") as env:
+        # Reset — get the task
+        result = await env.reset()
+        state = await env.state()
+        print(f"Task: {state.original_prompt}")
+        print(f"Difficulty: {state.task_difficulty}")
+
+        # Agent submits an optimized prompt
+        result = await env.step(
+            TokenOptimiserAction(optimized_prompt="Explain machine learning briefly.")
+        )
+        print(f"LLM Response: {result.observation.llm_response}")
+        print(f"Reward: {result.reward}")
+        print(f"Tokens — in: {result.observation.input_tokens}, out: {result.observation.output_tokens}")
+
+asyncio.run(main())
+```
+
+---
+
+## ☁️ Deployment
+
+Deploy to Hugging Face Spaces using the OpenEnv CLI:
+
+```bash
+openenv push
+```
+
+---
+
+## 🔧 Environment Variables Reference
+
+| Variable | Default | Description |
+|----------|---------|-------------|
+| `HF_TOKEN` | *(required)* | Hugging Face API key with Inference Providers permission |
+| `API_BASE_URL` | `https://router.huggingface.co/v1` | LLM endpoint URL |
+| `MODEL_NAME` | `Qwen/Qwen2.5-72B-Instruct` | Model used for responses and judging |
+| `SERVER_URL` | `http://localhost:8000` | Environment server URL (for inference.py) |
+| `LOCAL_IMAGE_NAME` | *(optional)* | Docker image name — auto-spins container if set |
+
+---
+
+## 📦 Dependencies
+
+- [`openenv-core`](https://github.com/meta-pytorch/OpenEnv) ≥ 0.2.2 — RL environment framework
+- `openai` — LLM API client (routed through HF)
+- `fastapi` + `uvicorn` — Environment server
+- `pydantic` v2 — Data model validation
+
+See [`pyproject.toml`](./pyproject.toml) for the full pinned dependency list.
+
+---
+
+## 📄 License
+
+BSD License — see source files for details.

__init__.py

@@ -0,0 +1,16 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Token Optimiser Environment."""
+
+from .client import TokenOptimiserEnv
+from .models import TokenOptimiserAction, TokenOptimiserObservation
+
+__all__ = [
+    "TokenOptimiserAction",
+    "TokenOptimiserObservation",
+    "TokenOptimiserEnv",
+]

client.py

@@ -0,0 +1,101 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# source tree.
+
+"""Token Optimiser Environment Client."""
+
+from typing import Dict
+
+from openenv.core import EnvClient
+from openenv.core.client_types import StepResult
+from openenv.core.env_server.types import State
+
+from .models import TokenOptimiserAction, TokenOptimiserObservation, TokenOptimiserState
+
+
+class TokenOptimiserEnv(
+    EnvClient[TokenOptimiserAction, TokenOptimiserObservation, TokenOptimiserState]
+):
+    """
+    Client for the Token Optimiser Environment.
+
+    This client maintains a persistent WebSocket connection to the environment server,
+    enabling efficient multi-step interactions with lower latency.
+    Each client instance has its own dedicated environment session on the server.
+
+    Example:
+        >>> # Connect to a running server
+        >>> with TokenOptimiserEnv(base_url="http://localhost:8000") as client:
+        ...     result = client.reset()
+        ...     print(result.observation.llm_response)
+        ...
+        ...     result = client.step(TokenOptimiserAction(optimized_prompt="Explain ML briefly"))
+        ...     print(result.observation.llm_response)
+
+    Example with Docker:
+        >>> # Automatically start container and connect
+        >>> client = TokenOptimiserEnv.from_docker_image("token_optimiser-env:latest")
+        >>> try:
+        ...     result = client.reset()
+        ...     result = client.step(TokenOptimiserAction(optimized_prompt="Test prompt"))
+        ... finally:
+        ...     client.close()
+    """
+
+    def _step_payload(self, action: TokenOptimiserAction) -> Dict:
+        """
+        Convert TokenOptimiserAction to JSON payload for step message.
+
+        Args:
+            action: TokenOptimiserAction instance
+
+        Returns:
+            Dictionary representation suitable for JSON encoding
+        """
+        return {
+            "optimized_prompt": action.optimized_prompt,
+        }
+
+    def _parse_result(self, payload: Dict) -> StepResult[TokenOptimiserObservation]:
+        """
+        Parse server response into StepResult[TokenOptimiserObservation].
+
+        Args:
+            payload: JSON response data from server
+
+        Returns:
+            StepResult with TokenOptimiserObservation
+        """
+        obs_data = payload.get("observation", {})
+        observation = TokenOptimiserObservation(
+            llm_response=obs_data.get("llm_response", ""),
+            input_tokens=obs_data.get("input_tokens", 0),
+            output_tokens=obs_data.get("output_tokens", 0),
+            reward=obs_data.get("reward", 0.0),
+        )
+
+        return StepResult(
+            observation=observation,
+            reward=payload.get("reward", 0.0),
+            done=payload.get("done", False),
+        )
+
+    def _parse_state(self, payload: Dict) -> TokenOptimiserState:
+        """
+        Parse server response into TokenOptimiserState object.
+
+        Args:
+            payload: JSON response from state request
+
+        Returns:
+            TokenOptimiserState object with episode information
+        """
+        return TokenOptimiserState(
+            episode_id=payload.get("episode_id"),
+            step_count=payload.get("step_count", 0),
+            original_prompt=payload.get("original_prompt", ""),
+            task_difficulty=payload.get("task_difficulty", "easy"),
+            task_index=payload.get("task_index", 0)
+        )

inference.py

@@ -0,0 +1,257 @@
+"""
+Inference Script — Token Optimiser Environment
+================================================
+STDOUT FORMAT (mandatory):
+  [START] task=<task> env=<env> model=<model>
+  [STEP]  step=<n> action=<str> reward=<0.00> done=<true|false> error=<msg|null>
+  [END]   success=<true|false> steps=<n> score=<0.000> rewards=<r1,r2,...>
+
+Environment variables required:
+  HF_TOKEN       — Hugging Face API key
+  API_BASE_URL   — LLM endpoint  (default: https://router.huggingface.co/v1)
+  MODEL_NAME     — Model id       (default: Qwen/Qwen2.5-72B-Instruct)
+  SERVER_URL     — Running env server (default: http://localhost:8000)
+  LOCAL_IMAGE_NAME — Docker image name (optional; spins up container if set)
+"""
+
+import asyncio
+import os
+import textwrap
+from typing import List, Optional
+
+from openai import OpenAI
+
+from token_optimiser import TokenOptimiserEnv, TokenOptimiserAction
+
+# ---------------------------------------------------------------------------
+# Configuration
+# ---------------------------------------------------------------------------
+API_BASE_URL: str = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")
+MODEL_NAME: str = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-72B-Instruct")
+HF_TOKEN: Optional[str] = os.getenv("HF_TOKEN") or os.getenv("API_KEY")
+SERVER_URL: str = os.getenv("SERVER_URL", "http://localhost:8000")
+LOCAL_IMAGE_NAME: Optional[str] = os.getenv("LOCAL_IMAGE_NAME")
+
+TASK_NAME: str = "token_optimization"
+BENCHMARK: str = "token_optimiser"
+MAX_STEPS: int = 5
+TEMPERATURE: float = 0.3
+MAX_TOKENS: int = 200
+SUCCESS_THRESHOLD: float = 0.6
+
+SYSTEM_PROMPT = textwrap.dedent("""
+    You are a prompt optimization expert. Rewrite the given prompt to:
+    1. Use the fewest possible tokens (concise language, no filler words)
+    2. Preserve full semantic meaning and intent
+    3. Add explicit output-format constraints (e.g., "in 5 bullet points", "as JSON with keys: …")
+    4. Guide the responder toward a shorter, precise answer
+
+    Reply with ONLY the optimized prompt — no explanations, no prefixes, no quotes.
+""").strip()
+
+# ---------------------------------------------------------------------------
+# Logging helpers (mandatory format)
+# ---------------------------------------------------------------------------
+
+def log_start(task: str, env: str, model: str) -> None:
+    print(f"[START] task={task} env={env} model={model}", flush=True)
+
+
+def log_step(step: int, action: str, reward: float, done: bool, error: Optional[str]) -> None:
+    # Truncate action for readability but keep it on one line
+    action_short = action.replace("\n", " ")[:120]
+    error_val = error if error else "null"
+    print(
+        f"[STEP] step={step} action={action_short!r} "
+        f"reward={reward:.2f} done={str(done).lower()} error={error_val}",
+        flush=True,
+    )
+
+
+def log_end(success: bool, steps: int, score: float, rewards: List[float]) -> None:
+    rewards_str = ",".join(f"{r:.2f}" for r in rewards)
+    print(
+        f"[END] success={str(success).lower()} steps={steps} "
+        f"score={score:.3f} rewards={rewards_str}",
+        flush=True,
+    )
+
+
+# ---------------------------------------------------------------------------
+# LLM helpers
+# ---------------------------------------------------------------------------
+
+def _build_user_message(original_prompt: str, step: int,
+                         prev_reward: float, prev_response: str,
+                         history: List[str]) -> str:
+    if step == 1:
+        return (
+            f"Optimize this prompt to minimize tokens while preserving all meaning:\n\n"
+            f"{original_prompt}"
+        )
+    history_block = "\n".join(history[-3:]) if history else "None"
+    return textwrap.dedent(f"""
+        Original prompt:
+        {original_prompt}
+
+        Your last optimized version got reward: {prev_reward:.2f}
+        LLM responded with: {prev_response!r}
+
+        Recent history:
+        {history_block}
+
+        Improve your optimization further. Reply with ONLY the new optimized prompt.
+    """).strip()
+
+
+def get_optimized_prompt(
+    llm: OpenAI,
+    original_prompt: str,
+    step: int,
+    prev_reward: float,
+    prev_response: str,
+    history: List[str],
+) -> str:
+    user_msg = _build_user_message(original_prompt, step, prev_reward, prev_response, history)
+    try:
+        completion = llm.chat.completions.create(
+            model=MODEL_NAME,
+            messages=[
+                {"role": "system", "content": SYSTEM_PROMPT},
+                {"role": "user", "content": user_msg},
+            ],
+            temperature=TEMPERATURE,
+            max_tokens=MAX_TOKENS,
+        )
+        result = (completion.choices[0].message.content or "").strip()
+        return result if result else "Explain briefly."
+    except Exception as exc:
+        print(f"[DEBUG] LLM call failed: {exc}", flush=True)
+        return _rule_based_compress(original_prompt, step)
+
+
+# Rule-based fallback compressor (used when LLM is unavailable)
+_FILLER = {
+    "please", "kindly", "could", "you", "can", "i", "need", "want", "would",
+    "like", "very", "really", "just", "actually", "basically", "specifically",
+    "a", "an", "the", "in", "of", "to", "and", "that", "is", "are", "be",
+    "will", "should", "must", "have", "has", "do", "does", "for", "with",
+    "as", "at", "by", "on", "or", "but", "it", "its", "this",
+}
+_BREVITY = [
+    "",                                    # step 1 — just strip fillers
+    " Be brief.",                           # step 2
+    " Limit response to 3 sentences.",      # step 3
+    " Reply in one sentence.",              # step 4+
+]
+
+
+def _rule_based_compress(original_prompt: str, step: int = 1) -> str:
+    """Strip filler words and add a conciseness constraint."""
+    words = original_prompt.split()
+    compressed = [
+        w for w in words
+        if w.lower().rstrip(".,?!") not in _FILLER
+    ]
+    suffix = _BREVITY[min(step - 1, len(_BREVITY) - 1)]
+    result = " ".join(compressed) + suffix
+    return result if result.strip() else original_prompt
+
+
+# ---------------------------------------------------------------------------
+# Main episode loop
+# ---------------------------------------------------------------------------
+
+async def run_episode(llm: OpenAI) -> None:
+    rewards: List[float] = []
+    steps_taken = 0
+    score = 0.0
+    success = False
+
+    log_start(task=TASK_NAME, env=BENCHMARK, model=MODEL_NAME)
+
+    # Connect to environment
+    if LOCAL_IMAGE_NAME:
+        env = await TokenOptimiserEnv.from_docker_image(LOCAL_IMAGE_NAME)
+    else:
+        env = TokenOptimiserEnv(base_url=SERVER_URL)
+        await env.connect()
+
+    try:
+        # Reset — get initial observation
+        reset_result = await env.reset()
+
+        # Fetch original prompt from server state
+        env_state = await env.state()
+        original_prompt: str = env_state.original_prompt or "Explain machine learning briefly."
+
+        print(
+            f"[DEBUG] Task difficulty={env_state.task_difficulty} | "
+            f"Original prompt ({len(original_prompt.split())} words): {original_prompt[:80]}...",
+            flush=True,
+        )
+
+        prev_reward = 0.0
+        prev_response = ""
+        history: List[str] = []
+
+        for step in range(1, MAX_STEPS + 1):
+            # Ask LLM to optimize the prompt
+            optimized = get_optimized_prompt(
+                llm, original_prompt, step, prev_reward, prev_response, history
+            )
+
+            # Step the environment with the optimized prompt
+            error_msg: Optional[str] = None
+            reward = 0.0
+            done = False
+            try:
+                result = await env.step(TokenOptimiserAction(optimized_prompt=optimized))
+                obs = result.observation
+                reward = result.reward          # server puts reward at top-level, not inside obs
+                done = result.done or (step >= MAX_STEPS)
+                prev_response = obs.llm_response
+                print(
+                    f"[DEBUG] tokens in={obs.input_tokens} out={obs.output_tokens}",
+                    flush=True,
+                )
+            except Exception as exc:
+                error_msg = str(exc)
+                done = True
+
+            rewards.append(reward)
+            steps_taken = step
+            prev_reward = reward
+            history.append(f"step={step} prompt={optimized!r:.60} reward={reward:.2f}")
+
+            log_step(step=step, action=optimized, reward=reward, done=done, error=error_msg)
+
+            if done:
+                break
+
+        # Score = average reward across steps, clamped to [0, 1]
+        score = sum(rewards) / len(rewards) if rewards else 0.0
+        score = max(0.0, min(1.0, score))
+        success = score >= SUCCESS_THRESHOLD
+
+    except Exception as exc:
+        print(f"[DEBUG] Episode error: {exc}", flush=True)
+    finally:
+        try:
+            await env.close()
+        except Exception:
+            pass
+        log_end(success=success, steps=steps_taken, score=score, rewards=rewards)
+
+
+async def main() -> None:
+    if not HF_TOKEN:
+        print("[ERROR] HF_TOKEN environment variable not set. Exiting.", flush=True)
+        return
+
+    llm = OpenAI(base_url=API_BASE_URL, api_key=HF_TOKEN)
+    await run_episode(llm)
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

models.py

@@ -0,0 +1,35 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# source tree.
+
+"""
+Data models for the Prompt & Response Token Optimization Environment.
+"""
+
+from openenv.core.env_server.types import Action, Observation, State
+from pydantic import Field
+
+
+class TokenOptimiserAction(Action):
+    """Action for the Token Optimiser environment - the optimized prompt."""
+
+    optimized_prompt: str = Field(..., description="The optimized prompt sent to the LLM")
+
+
+class TokenOptimiserObservation(Observation):
+    """Observation from the Token Optimiser environment - LLM response and token metrics."""
+
+    llm_response: str = Field(default="", description="The response from the LLM")
+    input_tokens: int = Field(default=0, description="Number of tokens in the optimized prompt")
+    output_tokens: int = Field(default=0, description="Number of tokens in the LLM response")
+    reward: float = Field(default=0.0, description="Reward score for this step (0.0-1.0)")
+
+
+class TokenOptimiserState(State):
+    """State for the Token Optimiser environment."""
+
+    original_prompt: str = Field(default="", description="The original user prompt/task")
+    task_difficulty: str = Field(default="easy", description="Current task difficulty: easy/medium/hard")
+    task_index: int = Field(default=0, description="Index of current task in task bank")

openenv.yaml

@@ -0,0 +1,7 @@
+spec_version: 1
+name: token_optimiser
+type: space
+runtime: fastapi
+app: server.app:app
+port: 8000
+

openenv_token_optimiser.egg-info/PKG-INFO

@@ -0,0 +1,9 @@
+Metadata-Version: 2.4
+Name: openenv-token_optimiser
+Version: 0.1.0
+Summary: Token Optimiser environment for OpenEnv
+Requires-Python: >=3.10
+Requires-Dist: openenv-core[core]>=0.2.2
+Provides-Extra: dev
+Requires-Dist: pytest>=8.0.0; extra == "dev"
+Requires-Dist: pytest-cov>=4.0.0; extra == "dev"

openenv_token_optimiser.egg-info/SOURCES.txt

@@ -0,0 +1,15 @@
+README.md
+pyproject.toml
+./__init__.py
+./client.py
+./inference.py
+./models.py
+openenv_token_optimiser.egg-info/PKG-INFO
+openenv_token_optimiser.egg-info/SOURCES.txt
+openenv_token_optimiser.egg-info/dependency_links.txt
+openenv_token_optimiser.egg-info/entry_points.txt
+openenv_token_optimiser.egg-info/requires.txt
+openenv_token_optimiser.egg-info/top_level.txt
+server/__init__.py
+server/app.py
+server/token_optimiser_environment.py

openenv_token_optimiser.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

openenv_token_optimiser.egg-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+server = token_optimiser.server.app:main

openenv_token_optimiser.egg-info/requires.txt

@@ -0,0 +1,5 @@
+openenv-core[core]>=0.2.2
+
+[dev]
+pytest>=8.0.0
+pytest-cov>=4.0.0

openenv_token_optimiser.egg-info/top_level.txt

@@ -0,0 +1 @@
+token_optimiser

pyproject.toml

@@ -0,0 +1,45 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+[build-system]
+requires = ["setuptools>=45", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "openenv-token_optimiser"
+version = "0.1.0"
+description = "Token Optimiser environment for OpenEnv"
+requires-python = ">=3.10"
+dependencies = [
+    # Core OpenEnv runtime (provides FastAPI server + HTTP client types)
+    # install from github
+    # "openenv-core[core] @ git+https://github.com/meta-pytorch/OpenEnv.git",
+    "openenv-core[core]>=0.2.2",
+    # Environment-specific dependencies
+    # Add all dependencies needed for your environment here
+    # Examples:
+    # "numpy>=1.19.0",
+    # "torch>=2.0.0",
+    # "gymnasium>=0.29.0",
+    # "openspiel>=1.0.0",
+    # "smolagents>=1.22.0,<2",
+]
+
+[project.optional-dependencies]
+dev = [
+    "pytest>=8.0.0",
+    "pytest-cov>=4.0.0",
+]
+
+[project.scripts]
+# Server entry point - enables running via: uv run --project . server
+# or: python -m token_optimiser.server.app
+server = "token_optimiser.server.app:main"
+
+[tool.setuptools]
+include-package-data = true
+packages = ["token_optimiser", "token_optimiser.server"]
+package-dir = { "token_optimiser" = ".", "token_optimiser.server" = "server" }

server/__init__.py

@@ -0,0 +1,11 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Token Optimiser environment server components."""
+
+from .token_optimiser_environment import TokenOptimiserEnvironment
+
+__all__ = ["TokenOptimiserEnvironment"]

server/app.py

@@ -0,0 +1,85 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# source tree.
+
+"""
+FastAPI application for the Token Optimiser Environment.
+
+This module creates an HTTP server that exposes the TokenOptimiserEnvironment
+over HTTP and WebSocket endpoints, compatible with EnvClient.
+
+Endpoints:
+    - POST /reset: Reset the environment
+    - POST /step: Execute an action
+    - GET /state: Get current environment state
+    - GET /schema: Get action/observation schemas
+    - WS /ws: WebSocket endpoint for persistent sessions
+
+Usage:
+    # Development (with auto-reload):
+    uvicorn server.app:app --reload --host 0.0.0.0 --port 8000
+
+    # Production:
+    uvicorn server.app:app --host 0.0.0.0 --port 8000 --workers 4
+
+    # Or run directly:
+    python -m server.app
+"""
+
+try:
+    from openenv.core.env_server.http_server import create_app
+except Exception as e:  # pragma: no cover
+    raise ImportError(
+        "openenv is required for the web interface. Install dependencies with '\n    uv sync\n'"
+    ) from e
+
+# Import using absolute paths from package root
+import sys
+import os
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))
+
+from models import TokenOptimiserAction, TokenOptimiserObservation, TokenOptimiserState
+from server.token_optimiser_environment import TokenOptimiserEnvironment
+
+
+# Create the app with web interface and README integration
+app = create_app(
+    TokenOptimiserEnvironment,
+    TokenOptimiserAction,
+    TokenOptimiserObservation,
+    env_name="token_optimiser",
+    max_concurrent_envs=10,  # Increased for RL training concurrency
+)
+
+
+def main(host: str = "0.0.0.0", port: int = 8000):
+    """
+    Entry point for direct execution via uv run or python -m.
+
+    This function enables running the server without Docker:
+        uv run --project . server
+        uv run --project . server --port 8001
+        python -m token_optimiser.server.app
+
+    Args:
+        host: Host address to bind to (default: "0.0.0.0")
+        port: Port number to listen on (default: 8000)
+
+    For production deployments, consider using uvicorn directly with
+    multiple workers:
+        uvicorn token_optimiser.server.app:app --workers 4
+    """
+    import uvicorn
+
+    uvicorn.run(app, host=host, port=port)
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--port", type=int, default=8000)
+    args = parser.parse_args()
+    main(port=args.port)

server/requirements.txt

@@ -0,0 +1,6 @@
+openenv[core]>=0.2.0
+fastapi>=0.115.0
+uvicorn>=0.24.0
+
+
+

server/token_optimiser_environment.py

@@ -0,0 +1,366 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# source tree.
+
+"""
+Prompt & Response Token Optimization Environment Implementation.
+
+A sandboxed LLM interaction environment where an AI agent optimizes both input prompts
+and expected output responses to minimize total token usage while maintaining correctness.
+"""
+
+import os
+import random
+from uuid import uuid4
+
+try:
+    from openai import OpenAI
+except ImportError:
+    OpenAI = None
+
+from openenv.core.env_server.interfaces import Environment
+from openenv.core.env_server.types import State
+
+try:
+    from ..models import TokenOptimiserAction, TokenOptimiserObservation, TokenOptimiserState
+except ImportError:
+    from models import TokenOptimiserAction, TokenOptimiserObservation, TokenOptimiserState
+
+
+class TokenOptimiserEnvironment(Environment):
+    """
+    Prompt & Response Token Optimization Environment.
+
+    The agent receives a user prompt/task and must optimize it to reduce token usage
+    while guiding the LLM to produce correct, properly formatted responses.
+    """
+
+    # Enable concurrent WebSocket sessions - REQUIRED for RL training
+    SUPPORTS_CONCURRENT_SESSIONS: bool = True
+
+    def __init__(self):
+        """Initialize the token optimization environment."""
+        self._state = TokenOptimiserState(episode_id=str(uuid4()), step_count=0)
+        self._task_bank = self._load_task_bank()
+        self._current_task = None
+        self._reset_count = 0
+
+        # Hybrid LLM client — reads credentials from env vars at startup
+        api_key = os.getenv("HF_TOKEN") or os.getenv("API_KEY")
+        api_base = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")
+        self._model = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-72B-Instruct")
+        if OpenAI and api_key:
+            self._llm = OpenAI(base_url=api_base, api_key=api_key)
+        else:
+            self._llm = None
+
+    def _load_task_bank(self):
+        """Load the bank of prompt optimization tasks."""
+        return [
+            # EASY TASK
+            {
+                "difficulty": "easy",
+                "prompt": "Can you please explain in a very detailed manner what machine learning is and how it works step by step?",
+                "expected_format": "brief explanation",
+                "reference_response": "Machine learning is a subset of AI that enables computers to learn from data without explicit programming. It works by identifying patterns in training data to make predictions or decisions on new data.",
+                "max_output_tokens": 50,
+                "description": "Reduce verbosity while preserving core concept"
+            },
+            # MEDIUM TASK
+            {
+                "difficulty": "medium",
+                "prompt": "I need a comprehensive analysis of the renewable energy market trends over the past decade, including solar, wind, and hydroelectric power growth rates, investment patterns, technological advancements, and policy impacts across different regions globally.",
+                "expected_format": "5 bullet points summarizing key trends",
+                "reference_response": "• Solar power capacity grew 22% annually avg.  • Wind energy investments reached $140B in 2020  • Hydroelectric remains largest renewable source  • Battery storage tech advancing rapidly  • Policy incentives driving global adoption",
+                "max_output_tokens": 100,
+                "description": "Compress input + specify bullet point format + length limit"
+            },
+            # HARD TASK
+            {
+                "difficulty": "hard",
+                "prompt": "As a senior data scientist, I need you to analyze our Q3 sales performance dataset and provide actionable insights. The dataset contains: customer demographics, purchase history, product categories, regional sales data, marketing campaign ROI, seasonal trends, and competitor analysis. Please identify: 1) Our top 3 performing product categories and why, 2) Geographic regions with highest growth potential, 3) Customer segments most responsive to our email campaigns, 4) Optimal marketing budget allocation for Q4, and 5) Risks to watch based on economic indicators.",
+                "expected_format": "JSON with 5 keys: top_categories, growth_regions, responsive_segments, budget_allocation, risks_watch",
+                "reference_response": '{"top_categories": ["electronics", "software", "home_goods"], "growth_regions": ["SE Asia", "Latin America", "Africa"], "responsive_segments": ["young_professionals", "tech_enthusiasts"], "budget_allocation": {"email": 0.3, "social": 0.25, "search": 0.2, "tv": 0.15, "other": 0.1}, "risks_watch": ["inflation", "supply_chain", "labor_shortage"]}',
+                "max_output_tokens": 200,
+                "description": "Multi-intent optimization: compress complex request + specify JSON format + accuracy + length constraints"
+            }
+        ]
+
+    def reset(self) -> TokenOptimiserObservation:
+        """
+        Reset the environment with a random task from the task bank.
+
+        Returns:
+            TokenOptimiserObservation with initial state
+        """
+        # Select a random task
+        self._current_task = random.choice(self._task_bank)
+        self._state = TokenOptimiserState(
+            episode_id=str(uuid4()),
+            step_count=0,
+            original_prompt=self._current_task["prompt"],
+            task_difficulty=self._current_task["difficulty"],
+            task_index=self._task_bank.index(self._current_task)
+        )
+        self._reset_count += 1
+
+        return TokenOptimiserObservation(
+            llm_response="",
+            input_tokens=0,
+            output_tokens=0,
+            reward=0.0
+        )
+
+    def step(self, action: TokenOptimiserAction) -> TokenOptimiserObservation:  # type: ignore[override]
+        """
+        Execute a step in the environment by evaluating the agent's optimized prompt.
+
+        Args:
+            action: TokenOptimiserAction containing the optimized prompt
+
+        Returns:
+            TokenOptimiserObservation with LLM response simulation and reward
+        """
+        self._state.step_count += 1
+        optimized_prompt = action.optimized_prompt
+        original_prompt = self._current_task["prompt"]
+
+        # 1. Call real LLM (or fallback) to get the actual response + token counts
+        llm_response, input_tokens, output_tokens = self._call_llm(optimized_prompt)
+
+        # 2. LLM-as-judge: semantic quality score (0.0-1.0)
+        semantic_score = self._judge_semantic_quality(original_prompt, llm_response)
+
+        # 3. Token efficiency: how much did we reduce vs the original prompt token count?
+        original_tokens = len(original_prompt.split()) * 1.3
+        ref_output_tokens = len(self._current_task["reference_response"].split()) * 1.3
+        ref_total = original_tokens + ref_output_tokens
+        actual_total = input_tokens + output_tokens
+        token_efficiency = max(0.0, min(0.4, (ref_total - actual_total) / max(ref_total, 1)))
+
+        # 4. Format compliance (0.0-0.2)
+        expected_fmt = self._current_task["expected_format"]
+        format_score = 0.0
+        if "bullet" in expected_fmt and any(c in llm_response for c in ("•", "*", "-", "\n")):
+            format_score = 0.2
+        elif "json" in expected_fmt.lower() and "{" in llm_response and "}" in llm_response:
+            format_score = 0.2
+        elif "brief" in expected_fmt and len(llm_response.split()) < 30:
+            format_score = 0.2
+
+        # 5. Length penalty if output way too long
+        max_out = self._current_task["max_output_tokens"]
+        length_penalty = -0.1 if output_tokens > max_out * 2 else 0.0
+
+        # Final reward: weighted hybrid
+        reward = (
+            token_efficiency            # 0.0 - 0.4   (token saving)
+            + semantic_score * 0.3      # 0.0 - 0.3   (LLM judge quality)
+            + format_score              # 0.0 - 0.2   (format compliance)
+            + length_penalty            # 0.0 or -0.1 (penalty)
+        )
+        reward = max(0.0, min(1.0, reward))
+
+        print(
+            f"[ENV] tok_eff={token_efficiency:.2f} semantic={semantic_score:.2f} "
+            f"fmt={format_score:.2f} => reward={reward:.2f}",
+            flush=True,
+        )
+
+        return TokenOptimiserObservation(
+            llm_response=llm_response,
+            input_tokens=int(input_tokens),
+            output_tokens=int(output_tokens),
+            reward=reward
+        )
+
+    def _call_llm(self, prompt: str) -> tuple[str, int, int]:
+        """
+        Call the real LLM with the optimized prompt.
+        Returns (response_text, input_tokens, output_tokens).
+        Falls back to rule-based simulation if LLM is unavailable.
+        """
+        if self._llm is not None:
+            try:
+                resp = self._llm.chat.completions.create(
+                    model=self._model,
+                    messages=[{"role": "user", "content": prompt}],
+                    max_tokens=200,
+                    temperature=0.3,
+                )
+                text = (resp.choices[0].message.content or "").strip()
+                in_tok = resp.usage.prompt_tokens if resp.usage else len(prompt.split())
+                out_tok = resp.usage.completion_tokens if resp.usage else len(text.split())
+                return text, in_tok, out_tok
+            except Exception as e:
+                print(f"[ENV] LLM call failed, using fallback: {e}")
+
+        # Rule-based fallback
+        return self._fallback_simulate(prompt)
+
+    def _fallback_simulate(self, prompt: str) -> tuple[str, int, int]:
+        """Fast deterministic fallback when LLM is unavailable."""
+        if self._current_task is None:
+            text = "No task loaded."
+            return text, len(prompt.split()), len(text.split())
+
+        expected_format = self._current_task["expected_format"]
+        original_words = len(self._current_task["prompt"].split())
+        compression_ratio = len(prompt.split()) / max(original_words, 1)
+
+        if "brief explanation" in expected_format:
+            text = ("Machine learning is AI that learns from data to make predictions."
+                    if compression_ratio <= 0.6
+                    else "Machine learning enables systems to learn from experience and improve without explicit programming.")
+        elif "bullet" in expected_format:
+            text = ("• Solar power growing rapidly\n• Wind energy investments increasing\n• Hydroelectric power stable\n• Battery storage advancing\n• Global renewable adoption rising"
+                    if compression_ratio <= 0.7
+                    else "Renewable energy sectors are growing, led by solar and wind with strong policy support.")
+        elif "json" in expected_format.lower():
+            text = ('{"top_categories": ["electronics", "software"], "growth_regions": ["Asia", "Africa"], "responsive_segments": ["professionals"], "budget_allocation": {"email": 0.4, "social": 0.3}, "risks_watch": ["inflation"]}'  # noqa
+                    if compression_ratio <= 0.6
+                    else "Key categories: electronics, software. Growth in Asia and Africa.")
+        else:
+            text = "I understand your request and will provide a helpful response."
+
+        in_tok = int(len(prompt.split()) * 1.3)
+        out_tok = int(len(text.split()) * 1.3)
+        return text, in_tok, out_tok
+
+    def _judge_semantic_quality(self, original_prompt: str, response: str) -> float:
+        """
+        LLM-as-judge: score how well the response answers the original prompt.
+        Returns a float 0.0-1.0.
+        """
+        if self._llm is None:
+            return self._keyword_fallback_score(original_prompt, response)
+
+        judge_prompt = (
+            f"Rate 0 to 10 how well the RESPONSE answers the ORIGINAL question. "
+            f"Consider accuracy and completeness. Reply with a single integer only.\n\n"
+            f"ORIGINAL: {original_prompt[:300]}\n\nRESPONSE: {response[:400]}"
+        )
+        try:
+            resp = self._llm.chat.completions.create(
+                model=self._model,
+                messages=[{"role": "user", "content": judge_prompt}],
+                max_tokens=5,
+                temperature=0.0,
+            )
+            raw = (resp.choices[0].message.content or "5").strip()
+            score = int("".join(c for c in raw if c.isdigit())[:2] or "5")
+            return min(max(score / 10.0, 0.0), 1.0)
+        except Exception:
+            return self._keyword_fallback_score(original_prompt, response)
+
+    def _keyword_fallback_score(self, original_prompt: str, response: str) -> float:
+        """Simple keyword overlap as semantic score when judge is unavailable."""
+        key_concepts = {
+            "machine", "learning", "ai", "data", "predict", "solar", "wind",
+            "energy", "renewable", "sales", "customer", "product", "market",
+            "budget", "analysis", "trend", "growth", "json", "bullet",
+        }
+        orig = set(original_prompt.lower().split()) & key_concepts
+        resp = set(response.lower().split()) & key_concepts
+        raw = (len(resp) / len(orig)) if orig else 0.5
+        return min(raw, 1.0)
+
+    def _calculate_reward(self, original_prompt: str, optimized_prompt: str,
+                         llm_response: str, expected_format: str, reference_response: str,
+                         input_tokens: int, output_tokens: int) -> float:
+        """
+        Calculate multi-component reward score (0.0-1.0).
+        """
+        # Component 1: Token Efficiency (0.0-0.4)
+        original_tokens = len(original_prompt.split()) * 1.3
+        optimized_input_tokens = input_tokens
+        output_tokens_estimate = output_tokens
+
+        # Reference token counts for comparison
+        ref_input_tokens = len(self._current_task["prompt"].split()) * 1.3
+        ref_output_tokens = len(self._current_task["reference_response"].split()) * 1.3
+        ref_total_tokens = ref_input_tokens + ref_output_tokens
+
+        actual_total_tokens = optimized_input_tokens + output_tokens_estimate
+        token_efficiency = max(0, (ref_total_tokens - actual_total_tokens) / ref_total_tokens)
+        token_efficiency = min(token_efficiency, 0.4)  # Cap at 0.4
+
+        # Component 2: Semantic Preservation (0.0-0.3)
+        # Simple keyword-based similarity (in practice, would use embeddings)
+        original_keywords = set(original_prompt.lower().split())
+        response_keywords = set(llm_response.lower().split())
+
+        # Extract key concepts from original prompt
+        key_concepts = {"machine", "learning", "AI", "data", "predict", "solar", "wind",
+                       "energy", "renewable", "sales", "customer", "product", "market",
+                       "budget", "analysis", "trend", "growth", "json", "bullet", "point"}
+
+        original_key_concepts = original_keywords & key_concepts
+        response_key_concepts = response_keywords & key_concepts
+
+        if len(original_key_concepts) > 0:
+            semantic_similarity = len(response_key_concepts) / len(original_key_concepts)
+        else:
+            semantic_similarity = 0.5  # Neutral if no key concepts found
+
+        semantic_score = min(semantic_similarity, 0.3)  # Cap at 0.3
+
+        # Component 3: Format Compliance (0.0-0.2)
+        format_score = 0.0
+        if "bullet point" in expected_format.lower() and ("•" in llm_response or "*" in llm_response or "-" in llm_response):
+            format_score = 0.2
+        elif "json" in expected_format.lower() and ("{" in llm_response and "}" in llm_response):
+            format_score = 0.2
+        elif "brief explanation" in expected_format.lower() and len(llm_response.split()) < 30:
+            format_score = 0.2
+
+        # Component 4: Length Appropriateness (0.0-0.1)
+        length_score = 0.0
+        max_expected = self._current_task["max_output_tokens"]
+        if output_tokens <= max_expected:
+            length_score = 0.1
+        elif output_tokens <= max_expected * 1.5:  # Partial credit
+            length_score = 0.05
+
+        # Component 5: Cost Simulation Bonus (0.0-0.05)
+        # Reward for being under reference token count
+        cost_bonus = 0.0
+        if actual_total_tokens < ref_total_tokens:
+            cost_bonus = min(0.05, (ref_total_tokens - actual_total_tokens) / ref_total_tokens * 0.05)
+
+        # Component 6: Latency Penalty (penalty)
+        latency_penalty = 0.0
+        if output_tokens > max_expected * 2:
+            latency_penalty = -0.1
+
+        # Component 7: Context Window Penalty (penalty)
+        context_penalty = 0.0
+        # Simulate context window limit (e.g., 4096 tokens)
+        if input_tokens > 3000:  # Assuming prompt + context
+            context_penalty = -0.1
+
+        # Calculate final reward
+        total_reward = (
+            token_efficiency +
+            semantic_score +
+            format_score +
+            length_score +
+            cost_bonus +
+            latency_penalty +
+            context_penalty
+        )
+
+        # Clamp to valid range
+        return max(0.0, min(1.0, total_reward))
+
+    @property
+    def state(self) -> TokenOptimiserState:
+        """
+        Get the current environment state.
+
+        Returns:
+            Current TokenOptimiserState
+        """
+        return self._state