yashu2000/TemporalBenchEnv

1# Copyright (c) Meta Platforms, Inc. and affiliates.2# All rights reserved.3#4# This source code is licensed under the BSD-style license found in the5# LICENSE file in the root directory of this source tree.6 7"""Typed OpenEnv client for TemporalBenchEnv."""8 9from typing import Any, Dict10 11from openenv.core.client_types import StepResult12from openenv.core.env_client import EnvClient13 14try:15    from env.models import (16        TemporalBenchAction,17        TemporalBenchObservation,18        TemporalBenchState,19    )20except ImportError:21    from TemporalBenchEnv.env.models import (22        TemporalBenchAction,23        TemporalBenchObservation,24        TemporalBenchState,25    )26 27 28class TemporalBenchEnvClient(29    EnvClient[30        TemporalBenchAction,31        TemporalBenchObservation,32        TemporalBenchState,33    ]34):35    """WebSocket client for TemporalBench MCQ episodes."""36 37    def _step_payload(self, action: TemporalBenchAction) -> Dict[str, Any]:38        payload: Dict[str, Any] = {"answer": action.answer}39        if action.confidence is not None:40            payload["confidence"] = action.confidence41        if action.reasoning is not None:42            payload["reasoning"] = action.reasoning43        return payload44 45    def _parse_result(self, payload: Dict[str, Any]) -> StepResult[TemporalBenchObservation]:46        obs_data = payload.get("observation")47        if not isinstance(obs_data, dict):48            obs_data = payload if isinstance(payload, dict) else {}49 50        done = payload.get("done", obs_data.get("done", False))51        reward = payload.get("reward", obs_data.get("reward"))52 53        observation = TemporalBenchObservation(54            step_idx=int(obs_data.get("step_idx", 0)),55            steps_remaining=int(obs_data.get("steps_remaining", 0)),56            max_steps=int(obs_data.get("max_steps", 9)),57            question=str(obs_data.get("question", "")),58            options=list(obs_data.get("options", [])),59            task_type=str(obs_data.get("task_type", "")),60            dataset=str(obs_data.get("dataset", "")),61            history=list(obs_data.get("history", [])),62            accuracy_so_far=float(obs_data.get("accuracy_so_far", 0.0)),63            done=done,64            reward=reward,65            metadata=obs_data.get("metadata", {}),66        )67        return StepResult(observation=observation, reward=reward, done=done)68 69    def _parse_state(self, payload: Dict[str, Any]) -> TemporalBenchState:70        state_data = payload.get("state")71        if not isinstance(state_data, dict):72            state_data = payload if isinstance(payload, dict) else {}73 74        return TemporalBenchState(75            episode_id=state_data.get("episode_id"),76            step_count=int(state_data.get("step_count", 0)),77            total_correct=int(state_data.get("total_correct", 0)),78            total_questions=int(state_data.get("total_questions", 9)),79            current_accuracy=float(state_data.get("current_accuracy", 0.0)),80            primary_domain=str(state_data.get("primary_domain", "PSML")),81            per_task_type_accuracy=dict(state_data.get("per_task_type_accuracy", {})),82            total_reward=float(state_data.get("total_reward", 0.0)),83        )84 85 86TemporalbenchenvEnv = TemporalBenchEnvClient87