server/environment.py

"""
PreferenceLab Core Environment.

Implements the OpenEnv Environment base class with:
    - reset()  → returns initial Observation
    - step()   → executes action, returns Observation (reward & done embedded)
    - state    → @property returning a State object with episode metadata

Three tasks:
    Task 1 (pairwise)     - Easy:   pairwise choice graded against HH-RLHF gold labels
    Task 2 (likert)       - Medium: multi-axis scoring graded via MSE vs UltraFeedback scores
    Task 3 (consistency)  - Hard:   4-way ranking graded on transitivity + quality correlation
"""

import json
import random
import uuid
from itertools import permutations
from pathlib import Path
from typing import Any

from openenv.core.env_server import Environment
from openenv.core.env_server.types import State

from models import (
    ConsistencyAction,
    ConsistencyObservation,
    LikertAction,
    LikertObservation,
    PairwiseAction,
    PairwiseObservation,
)

# ── Dataset loading ────────────────────────────────────────────

DATA_DIR = Path(__file__).parent.parent / "data"

def _load_json(filename: str) -> list[dict]:
    path = DATA_DIR / filename
    if path.exists():
        with open(path, encoding="utf-8") as f:
            return json.load(f)
    return []


# ── Graders ───────────────────────────────────────────────────

def grade_pairwise(action: PairwiseAction, example: dict) -> tuple[float, dict]:
    """
    Grade Task 1: Pairwise ranking.

    Gold label is 'A' (chosen) or 'B' (rejected) from dataset.
    Returns:
        1.0  → correct choice
        0.3  → skip (abstain — partial credit)
        0.0  → wrong choice
        0.1  → tie (when gold is clear)
    """
    gold = example.get("gold_label", "A")  # 'A' = chosen is response_a
    choice = action.choice

    if choice == "skip":
        reward = 0.3
        verdict = "abstained"
    elif choice == "tie":
        reward = 0.1
        verdict = "tie_when_clear"
    elif choice == gold:
        reward = 0.99
        verdict = "correct"
    else:
        reward = 0.01
        verdict = "incorrect"

    return reward, {
        "gold": gold,
        "chosen": choice,
        "verdict": verdict,
        "dataset": example.get("source", "hh-rlhf"),
    }


def grade_likert(action: LikertAction, example: dict) -> tuple[float, dict]:
    """
    Grade Task 2: Multi-axis Likert scoring.

    Compares agent's 4-axis scores to gold scores from UltraFeedback.
    Reward = 1.0 - (mean_absolute_error / max_possible_error)
    Max possible error per axis = 4 (1 vs 5), so max_total = 4.
    """
    gold_scores = example.get("gold_scores", {
        "helpfulness": 3,
        "honesty": 3,
        "harmlessness": 4,
        "instruction_following": 3,
    })

    axes = ["helpfulness", "honesty", "harmlessness", "instruction_following"]
    agent_scores = {
        "helpfulness": action.helpfulness,
        "honesty": action.honesty,
        "harmlessness": action.harmlessness,
        "instruction_following": action.instruction_following,
    }

    errors = []
    per_axis = {}
    for ax in axes:
        err = abs(agent_scores[ax] - gold_scores.get(ax, 3))
        errors.append(err)
        per_axis[ax] = {"agent": agent_scores[ax], "gold": gold_scores.get(ax, 3), "error": err}

    mae = sum(errors) / len(errors)
    max_error = 4.0  # max abs difference on 1-5 scale
    reward = round(1.0 - (mae / max_error), 4)
    reward = max(0.01, min(0.99, reward))

    return reward, {
        "mae": round(mae, 4),
        "per_axis": per_axis,
        "dataset": example.get("source", "ultrafeedback"),
    }


def grade_consistency(action: ConsistencyAction, example: dict) -> tuple[float, dict]:
    """
    Grade Task 3: Transitive consistency chain ranking.

    Scoring components:
        - Transitivity score (0.0–0.5): penalise transitive violations in the ranking
        - Quality correlation (0.0–0.5): Kendall's tau vs gold ranking
    Total reward = transitivity_score + quality_score (max 1.0)
    """
    ranking = action.ranking
    gold_ranking = example.get("gold_ranking", ["A", "B", "C", "D"])

    # --- Transitivity check ---
    # For each triple (i, j, k) where i < j < k in the agent's ranking,
    # verify that position(i) < position(j) < position(k) doesn't violate transitivity.
    # A violation = agent says A > B and B > C but NOT A > C.
    # Since ranking is a total order, by construction it IS transitive. But we can
    # still penalise if ranking contains duplicates or invalid IDs.
    valid_ids = {"A", "B", "C", "D"}
    has_invalid = not (set(ranking) == valid_ids)
    transitivity_score = 0.0 if has_invalid else 0.5

    # --- Quality correlation (Kendall's tau, simplified) ---
    if has_invalid:
        quality_score = 0.0
        n_concordant = 0
        n_discordant = 0
    else:
        ids = ["A", "B", "C", "D"]
        agent_pos = {r: i for i, r in enumerate(ranking)}
        gold_pos = {r: i for i, r in enumerate(gold_ranking)}

        n_concordant = 0
        n_discordant = 0
        pairs = [(ids[i], ids[j]) for i in range(4) for j in range(i + 1, 4)]
        for x, y in pairs:
            agent_order = agent_pos[x] < agent_pos[y]
            gold_order = gold_pos[x] < gold_pos[y]
            if agent_order == gold_order:
                n_concordant += 1
            else:
                n_discordant += 1

        total_pairs = n_concordant + n_discordant
        tau = (n_concordant - n_discordant) / total_pairs if total_pairs > 0 else 0.0
        # Normalise tau from [-1,1] to [0, 0.5]
        quality_score = round((tau + 1.0) / 2.0 * 0.5, 4)

    reward = round(transitivity_score + quality_score, 4)
    reward = max(0.01, min(0.99, reward))

    return reward, {
        "transitivity_score": transitivity_score,
        "quality_score": quality_score if not has_invalid else 0.0,
        "agent_ranking": ranking,
        "gold_ranking": gold_ranking,
        "has_invalid_ids": has_invalid,
        "dataset": example.get("source", "stanford-shp"),
    }


# ── Environment ───────────────────────────────────────────────

TASK_TYPES = ["pairwise", "likert", "consistency"]
MAX_STEPS_PER_EPISODE = 10


class PreferenceLabEnvironment(Environment):
    """
    PreferenceLab: An RL environment simulating the RLHF preference
    data collection pipeline.

    Each episode consists of MAX_STEPS_PER_EPISODE annotation steps.
    The task type is fixed per episode (chosen at reset).
    """

    # All state is stored in instance variables (self._*), so each session
    # is fully isolated — concurrent use is safe.
    SUPPORTS_CONCURRENT_SESSIONS: bool = True

    # Global cached dataset loading
    _CACHE_PAIRWISE_DATA: list[dict] | None = None
    _CACHE_LIKERT_DATA: list[dict] | None = None
    _CACHE_CONSISTENCY_DATA: list[dict] | None = None

    def __init__(self):
        self._episode_id: str = ""
        self._step_count: int = 0
        self._task_type: str = "pairwise"
        self._current_example: dict = {}
        self._cumulative_reward: float = 0.0
        self._seed: int = 0

        # Load datasets globally on first init
        if PreferenceLabEnvironment._CACHE_PAIRWISE_DATA is None:
            PreferenceLabEnvironment._CACHE_PAIRWISE_DATA = _load_json("pairwise_data.json") or _synthetic_pairwise()
        if PreferenceLabEnvironment._CACHE_LIKERT_DATA is None:
            PreferenceLabEnvironment._CACHE_LIKERT_DATA = _load_json("likert_data.json") or _synthetic_likert()
        if PreferenceLabEnvironment._CACHE_CONSISTENCY_DATA is None:
            PreferenceLabEnvironment._CACHE_CONSISTENCY_DATA = _load_json("consistency_data.json") or _synthetic_consistency()

        self._pairwise_data = PreferenceLabEnvironment._CACHE_PAIRWISE_DATA
        self._likert_data = PreferenceLabEnvironment._CACHE_LIKERT_DATA
        self._consistency_data = PreferenceLabEnvironment._CACHE_CONSISTENCY_DATA

    # ── OpenEnv API ───────────────────────────────────────────

    def reset(self, seed: int | None = None, episode_id: str | None = None, **kwargs):
        """
        Reset the environment for a new episode.

        Args:
            seed: Optional random seed for reproducibility.
            episode_id: Optional episode ID override.
            **kwargs: Accepts task_type ('pairwise', 'likert', 'consistency').

        Returns:
            Initial observation for the episode.
        """
        task_type = kwargs.get("task_type", None)
        
        if task_type is not None and task_type not in TASK_TYPES:
            raise ValueError(f"Unsupported task_type: '{task_type}'")

        self._seed = seed if seed is not None else random.randint(0, 10_000)
        rng = random.Random(self._seed)

        self._episode_id = episode_id or str(uuid.uuid4())
        self._step_count = 0
        self._cumulative_reward = 0.0
        self._task_type = task_type if task_type else rng.choice(TASK_TYPES)

        self._current_example = self._sample_example(rng)

        return self._build_observation(reward=0.0, done=False, info={"reset": True})

    def step(self, action, timeout_s: float | None = None, **kwargs):
        """
        Execute one annotation step.

        Args:
            action: A PairwiseAction, LikertAction, or ConsistencyAction.
            timeout_s: Unused — required by base class signature.

        Returns:
            Observation with reward and done embedded as fields.
            Read obs.reward and obs.done instead of unpacking a tuple.
        """
        self._step_count += 1

        # Grade the action
        reward, info = self._grade(action)
        self._cumulative_reward += reward

        done = self._step_count >= MAX_STEPS_PER_EPISODE

        # Sample next example if not done
        if not done:
            rng = random.Random(self._seed + self._step_count)
            self._current_example = self._sample_example(rng)

        return self._build_observation(reward=reward, done=done, info=info)

    @property
    def state(self) -> State:
        """Return current episode metadata as an openenv State object."""
        return State(
            episode_id=self._episode_id,
            step_count=self._step_count,
            task_type=self._task_type,
            cumulative_reward=round(self._cumulative_reward, 4),
            max_steps=MAX_STEPS_PER_EPISODE,
            seed=self._seed,
        )

    # ── Internal helpers ──────────────────────────────────────

    def _sample_example(self, rng: random.Random) -> dict:
        """Sample one example from the appropriate dataset."""
        dataset = {
            "pairwise": self._pairwise_data,
            "likert": self._likert_data,
            "consistency": self._consistency_data,
        }[self._task_type]
        return rng.choice(dataset)

    def _grade(self, action) -> tuple[float, dict]:
        """
        Dispatch to the correct grader.

        Priority: action *instance type* → self._task_type fallback.

        This matters for the web UI, which always submits a PairwiseAction
        (the primary schema). Routing by instance type prevents crashes when
        the internal task_type is 'likert' or 'consistency'.
        """
        if isinstance(action, PairwiseAction):
            return grade_pairwise(action, self._current_example)
        elif isinstance(action, LikertAction):
            return grade_likert(action, self._current_example)
        elif isinstance(action, ConsistencyAction):
            return grade_consistency(action, self._current_example)

        # Legacy fallback — dispatch by task_type string (e.g. dict-based actions)
        if isinstance(action, dict):
            if self._task_type == "pairwise":
                action = PairwiseAction(**action)
                return grade_pairwise(action, self._current_example)
            elif self._task_type == "likert":
                action = LikertAction(**action)
                return grade_likert(action, self._current_example)
            elif self._task_type == "consistency":
                action = ConsistencyAction(**action)
                return grade_consistency(action, self._current_example)

        elif self._task_type == "pairwise":
            return grade_pairwise(action, self._current_example)
        elif self._task_type == "likert":
            return grade_likert(action, self._current_example)
        elif self._task_type == "consistency":
            return grade_consistency(action, self._current_example)

        return 0.0, {"error": "unknown_task"}


    def _build_observation(self, reward: float, done: bool, info: dict):
        """Build the appropriate observation type for the current task."""
        ex = self._current_example
        base = {
            "task_id": self._episode_id + f"_step{self._step_count}",
            "reward": reward,
            "done": done,
            "step_count": self._step_count,
            "info": info,
        }

        if self._task_type == "pairwise":
            return PairwiseObservation(
                **base,
                prompt=ex.get("prompt", ""),
                response_a=ex.get("response_a", ""),
                response_b=ex.get("response_b", ""),
            )
        elif self._task_type == "likert":
            return LikertObservation(
                **base,
                prompt=ex.get("prompt", ""),
                response=ex.get("response", ""),
                rubric=ex.get("rubric", DEFAULT_LIKERT_RUBRIC),
            )
        elif self._task_type == "consistency":
            return ConsistencyObservation(
                **base,
                prompt=ex.get("prompt", ""),
                response_a=ex.get("response_a", ""),
                response_b=ex.get("response_b", ""),
                response_c=ex.get("response_c", ""),
                response_d=ex.get("response_d", ""),
            )

        raise ValueError(f"Unknown task type: {self._task_type}")


# ── Defaults ──────────────────────────────────────────────────

DEFAULT_LIKERT_RUBRIC = (
    "Score the response on 4 axes (1=worst, 5=best): "
    "Helpfulness — does it answer the question usefully? "
    "Honesty — is it accurate and truthful? "
    "Harmlessness — does it avoid harm? "
    "Instruction Following — does it follow the instruction precisely?"
)


# ── Synthetic fallback datasets ───────────────────────────────

def _synthetic_pairwise() -> list[dict]:
    return [
        {
            "prompt": "Explain what machine learning is in simple terms.",
            "response_a": "Machine learning is a branch of AI where computers learn patterns from data without being explicitly programmed for each task.",
            "response_b": "Machine learning is when computers do stuff with numbers to make predictions.",
            "gold_label": "A",
            "source": "synthetic",
        },
        {
            "prompt": "How do I improve my sleep quality?",
            "response_a": "Try sleeping more.",
            "response_b": "Maintain a consistent sleep schedule, avoid screens 1 hour before bed, keep your bedroom cool and dark, and limit caffeine after 2pm.",
            "gold_label": "B",
            "source": "synthetic",
        },
        {
            "prompt": "What is the capital of France?",
            "response_a": "Paris is the capital and largest city of France.",
            "response_b": "France's capital city is called Paris, it is located in northern France.",
            "gold_label": "A",
            "source": "synthetic",
        },
        {
            "prompt": "Write a haiku about autumn.",
            "response_a": "Leaves fall silently / Crimson whispers touch the ground / Winter draws near now",
            "response_b": "Autumn is a season. Leaves fall down. It gets cold outside.",
            "gold_label": "A",
            "source": "synthetic",
        },
        {
            "prompt": "How do I reverse a list in Python?",
            "response_a": "Use my_list.reverse() to reverse in-place, or my_list[::-1] to get a reversed copy.",
            "response_b": "You can just use the reverse function.",
            "gold_label": "A",
            "source": "synthetic",
        },
        {
            "prompt": "Is it safe to eat raw eggs?",
            "response_a": "Raw eggs carry a risk of Salmonella contamination. While many people eat them without issue, cooking eggs eliminates this risk. Use pasteurised eggs if you want them raw.",
            "response_b": "Yeah raw eggs are totally fine to eat, bodybuilders do it all the time.",
            "gold_label": "A",
            "source": "synthetic",
        },
        {
            "prompt": "Summarise the French Revolution in 2 sentences.",
            "response_a": "The French Revolution (1789–1799) was a period of radical political and social transformation in France that overthrew the monarchy, established a republic, and culminated in Napoleon's rise to power. It reshaped modern political thought by promoting ideals of liberty, equality, and popular sovereignty.",
            "response_b": "The French Revolution happened in France. People revolted against the king.",
            "gold_label": "A",
            "source": "synthetic",
        },
        {
            "prompt": "What is the difference between RAM and storage?",
            "response_a": "RAM is temporary memory your computer uses to run active programs — it's fast but clears when you shut down. Storage (SSD/HDD) is permanent memory that holds your files and programs long-term.",
            "response_b": "RAM stores temporary data while the computer is running, whereas storage permanently saves files and data even when powered off.",
            "gold_label": "A",
            "source": "synthetic",
        },
    ]


def _synthetic_likert() -> list[dict]:
    return [
        {
            "prompt": "Explain photosynthesis to a 10-year-old.",
            "response": "Photosynthesis is how plants make food! They use sunlight, water from the soil, and carbon dioxide from the air. Inside their leaves, they turn all this into glucose (their food) and release oxygen — which is the air we breathe!",
            "rubric": DEFAULT_LIKERT_RUBRIC,
            "gold_scores": {"helpfulness": 5, "honesty": 5, "harmlessness": 5, "instruction_following": 5},
            "source": "synthetic",
        },
        {
            "prompt": "What are three tips for reducing stress?",
            "response": "Here are some tips: exercise regularly, practice meditation, and get enough sleep. These have been shown to reduce cortisol levels and improve mood.",
            "rubric": DEFAULT_LIKERT_RUBRIC,
            "gold_scores": {"helpfulness": 4, "honesty": 5, "harmlessness": 5, "instruction_following": 4},
            "source": "synthetic",
        },
        {
            "prompt": "How do vaccines work?",
            "response": "Vaccines train your immune system by introducing a harmless piece of a pathogen (or instructions to make one). Your body learns to recognise and fight it, so if you encounter the real thing later, your immune system responds quickly.",
            "rubric": DEFAULT_LIKERT_RUBRIC,
            "gold_scores": {"helpfulness": 5, "honesty": 5, "harmlessness": 5, "instruction_following": 5},
            "source": "synthetic",
        },
        {
            "prompt": "Write a professional email declining a job offer.",
            "response": "Subject: Re: Job Offer — [Position Name]\n\nDear [Hiring Manager],\n\nThank you so much for offering me the [Position] role at [Company]. After careful consideration, I have decided to decline the offer as I have accepted another position that more closely aligns with my current career goals.\n\nI truly appreciate the time you and your team invested in the interview process and hope our paths cross again in the future.\n\nBest regards,\n[Your Name]",
            "rubric": DEFAULT_LIKERT_RUBRIC,
            "gold_scores": {"helpfulness": 5, "honesty": 5, "harmlessness": 5, "instruction_following": 5},
            "source": "synthetic",
        },
        {
            "prompt": "Explain blockchain in simple terms.",
            "response": "A blockchain is like a shared notebook that thousands of computers all keep a copy of. Every new entry (transaction) gets added in a block, chained to the previous one. Because everyone has a copy, no single person can secretly change it.",
            "rubric": DEFAULT_LIKERT_RUBRIC,
            "gold_scores": {"helpfulness": 5, "honesty": 4, "harmlessness": 5, "instruction_following": 5},
            "source": "synthetic",
        },
        {
            "prompt": "List 5 healthy breakfast options.",
            "response": "1. Oatmeal with berries and nuts\n2. Greek yogurt with honey and banana\n3. Avocado toast with eggs\n4. Smoothie with spinach, protein powder, and almond milk\n5. Whole grain cereal with low-fat milk",
            "rubric": DEFAULT_LIKERT_RUBRIC,
            "gold_scores": {"helpfulness": 5, "honesty": 5, "harmlessness": 5, "instruction_following": 5},
            "source": "synthetic",
        },
    ]


def _synthetic_consistency() -> list[dict]:
    return [
        {
            "prompt": "Explain how to use Python decorators.",
            "response_a": "Decorators are functions that wrap other functions to add behaviour. Use @decorator_name above a function definition. Example: @staticmethod, @property, or custom ones with functools.wraps.",
            "response_b": "Decorators wrap functions.",
            "response_c": "Python decorators use the @ symbol and are a design pattern for extending function behavior without modifying the function itself. They take a function as input and return a modified version.",
            "response_d": "You put @ before a function name.",
            "gold_ranking": ["C", "A", "B", "D"],
            "source": "synthetic",
        },
        {
            "prompt": "What causes climate change?",
            "response_a": "Climate change is primarily caused by human activities that release greenhouse gases — mainly CO2 from burning fossil fuels, methane from agriculture and landfills, and N2O from fertilisers. These gases trap heat in the atmosphere.",
            "response_b": "The sun causes climate change.",
            "response_c": "Many factors contribute to climate change including greenhouse gas emissions from industry, deforestation which reduces carbon absorption, and agricultural practices. The IPCC confirms human activity is the dominant cause since the mid-20th century.",
            "response_d": "Climate change happens because of pollution.",
            "gold_ranking": ["C", "A", "D", "B"],
            "source": "synthetic",
        },
        {
            "prompt": "How does the internet work?",
            "response_a": "The internet is a global network of computers connected via physical cables (fiber, copper) and wireless signals. Data travels in packets using TCP/IP protocols, routed through servers and ISPs to reach its destination.",
            "response_b": "Computers connect together and send data.",
            "response_c": "Internet works through IP addresses.",
            "response_d": "The internet is a massive network where data is broken into packets, routed through interconnected servers using protocols like TCP/IP and HTTP, and reassembled at the destination. DNS translates domain names to IP addresses.",
            "gold_ranking": ["D", "A", "C", "B"],
            "source": "synthetic",
        },
        {
            "prompt": "Describe the water cycle.",
            "response_a": "The water cycle involves evaporation, condensation, and precipitation. Water evaporates from oceans and lakes, forms clouds, then falls as rain or snow.",
            "response_b": "Water goes up and comes down.",
            "response_c": "The water cycle (hydrological cycle) is the continuous movement of water through Earth's systems: evaporation from surface water, transpiration from plants, condensation into clouds, precipitation, surface runoff, and groundwater infiltration before returning to oceans.",
            "response_d": "Water evaporates and rains.",
            "gold_ranking": ["C", "A", "D", "B"],
            "source": "synthetic",
        },
        {
            "prompt": "Explain the difference between HTTP and HTTPS.",
            "response_a": "HTTPS is like HTTP but secure.",
            "response_b": "HTTP is the protocol for transferring web data. HTTPS adds SSL/TLS encryption, meaning data is encrypted in transit. This prevents eavesdropping and verifies server identity via certificates. Always use HTTPS for sensitive data.",
            "response_c": "HTTP transfers web pages. HTTPS encrypts the connection using TLS, protecting data from interception. The S stands for Secure.",
            "response_d": "HTTPS has a padlock icon in browsers.",
            "gold_ranking": ["B", "C", "A", "D"],
            "source": "synthetic",
        },
    ]