permanence/openenv_env.py

"""
PERMANENCE — OpenEnv-compliant Environment subclass.

This module wraps the core ``PermanenceEnv`` (Gym-style) in an
``openenv.core.Environment`` subclass so the environment integrates
natively with the OpenEnv framework, ``create_fastapi_app``, TRL
rollout functions, and HuggingFace Spaces deployment.

The core logic (world state, actions, rewards) lives in the existing
``permanence/`` package and is untouched.  This file is pure adapter.
"""
from __future__ import annotations

import uuid
from typing import Any, Optional

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

from .env import PermanenceEnv
from .reward.rubrics import build_permanence_rubric

# Import from the top-level models module (sits next to server/, training/, etc.)
import sys, pathlib  # noqa: E401,E402
_project_root = str(pathlib.Path(__file__).resolve().parent.parent)
if _project_root not in sys.path:
    sys.path.insert(0, _project_root)

from models import PermanenceAction, PermanenceObservation, PermanenceState  # noqa: E402


class PermanenceOpenEnv(Environment[PermanenceAction, PermanenceObservation, PermanenceState]):
    """
    OpenEnv-native wrapper around the core PermanenceEnv.

    Implements the three abstract members required by
    ``openenv.core.Environment``:

    * ``reset(seed, episode_id, **kw) -> PermanenceObservation``
    * ``step(action, timeout_s, **kw) -> PermanenceObservation``
    * ``state`` property -> ``PermanenceState``
    """

    SUPPORTS_CONCURRENT_SESSIONS: bool = True

    def __init__(self) -> None:
        super().__init__()
        # Expose the composable rubric tree as the framework-standard
        # `rubric` attribute — used by tools like OpenEnv inspectors
        # and required by the hackathon grading criterion that explicitly
        # calls out composable-rubric usage.
        self.rubric = build_permanence_rubric()
        self._env: Optional[PermanenceEnv] = None
        self._episode_id: str = ""
        self._last_terminated: bool = False
        self._last_truncated: bool = False
        self._last_reason: Optional[str] = None

    # ------------------------------------------------------------------
    # reset
    # ------------------------------------------------------------------
    def reset(
        self,
        seed: Optional[int] = None,
        episode_id: Optional[str] = None,
        **kwargs: Any,
    ) -> PermanenceObservation:
        task_id = kwargs.get("task_id", None)
        difficulty = float(kwargs.get("difficulty", 0.5))
        config: Dict[str, Any] = {}
        if task_id:
            config["force_task"] = task_id
        self._env = PermanenceEnv(config=config)
        self._episode_id = episode_id or str(uuid.uuid4())[:8]
        self._last_terminated = False
        self._last_truncated = False
        self._last_reason = None

        obs_dict, info = self._env.reset(seed=seed, options={"difficulty": difficulty})

        return PermanenceObservation(
            text=obs_dict.get("text", ""),
            step=obs_dict.get("step", 0),
            task_id=obs_dict.get("task_id", ""),
            available_actions=obs_dict.get("available_actions", ""),
            done=False,
            reward=None,
            metadata=info,
        )

    # ------------------------------------------------------------------
    # step
    # ------------------------------------------------------------------
    def step(
        self,
        action: PermanenceAction,
        timeout_s: Optional[float] = None,
        **kwargs: Any,
    ) -> PermanenceObservation:
        # In HTTP mode, create_fastapi_app creates a fresh env per request.
        # Auto-reset if step is called on an uninitialised instance.
        if self._env is None:
            self.reset()

        obs_dict, reward, terminated, truncated, info = self._env.step(action.text)

        done = terminated or truncated
        self._last_terminated = terminated
        self._last_truncated = truncated
        self._last_reason = info.get("termination_reason")

        return PermanenceObservation(
            text=obs_dict.get("text", ""),
            step=obs_dict.get("step", 0),
            task_id=obs_dict.get("task_id", ""),
            available_actions=obs_dict.get("available_actions", ""),
            done=done,
            reward=float(reward) if done else None,
            metadata={
                **info,
                "episode_id": self._episode_id,
                "terminated": terminated,
                "truncated": truncated,
            },
        )

    # ------------------------------------------------------------------
    # state (property — required abstract)
    # ------------------------------------------------------------------
    @property
    def state(self) -> PermanenceState:
        if self._env is None or self._env._current_world_state is None:
            return PermanenceState(
                episode_id=self._episode_id or "not_started",
                step_count=0,
            )

        ws = self._env._current_world_state
        task = self._env._current_task

        return PermanenceState(
            episode_id=self._episode_id,
            step_count=self._env.episode_tracker.step_count,
            task_id=ws.task_id,
            task_difficulty=getattr(task, "difficulty", 0),
            locked_actions=sorted(ws.locked_actions.keys()),
            critical_options=dict(ws.critical_options),
            terminated=self._last_terminated,
            truncated=self._last_truncated,
            termination_reason=self._last_reason,
        )

    # ------------------------------------------------------------------
    # get_metadata (optional override for richer info)
    # ------------------------------------------------------------------
    def get_metadata(self) -> EnvironmentMetadata:
        return EnvironmentMetadata(
            name="PERMANENCE",
            description=(
                "First OpenEnv environment with persistent within-episode world state. "
                "Trains agents to predict action reversibility before acting."
            ),
            version="1.1.0",
            author="chanikya",
        )

    # ------------------------------------------------------------------
    # close
    # ------------------------------------------------------------------
    def close(self) -> None:
        self._env = None