Spaces:
Runtime error
Runtime error
| """Rubric-style task graders for validator compatibility.""" | |
| from __future__ import annotations | |
| from typing import Any | |
| try: | |
| from openenv.core.rubric import Rubric as OpenEnvRubric # type: ignore | |
| except Exception: | |
| try: | |
| from openenv.core.rubrics import Rubric as OpenEnvRubric # type: ignore | |
| except Exception: | |
| try: | |
| from openenv.rubric import Rubric as OpenEnvRubric # type: ignore | |
| except Exception: | |
| try: | |
| from openenv.rubrics import Rubric as OpenEnvRubric # type: ignore | |
| except Exception: | |
| class OpenEnvRubric: # pragma: no cover | |
| def reset(self) -> None: | |
| return None | |
| class _RubricMeta(type(OpenEnvRubric)): | |
| """Allow both RubricClass() and RubricClass(episode_data).""" | |
| def __call__(cls, *args: Any, **kwargs: Any) -> Any: | |
| if args and isinstance(args[0], dict) and len(args) == 1 and not kwargs: | |
| instance = super().__call__() | |
| return instance(args[0]) | |
| return super().__call__(*args, **kwargs) | |
| _MAX_REWARD_PER_QUERY: float = 1.5 | |
| _MAX_REWARDS: dict[str, float] = { | |
| "task_1": 1.0 * _MAX_REWARD_PER_QUERY, | |
| "task_2": 3.0 * _MAX_REWARD_PER_QUERY, | |
| "task_3": 5.0 * _MAX_REWARD_PER_QUERY, | |
| } | |
| _SCORE_MIN: float = 0.1 | |
| _SCORE_MAX: float = 0.99 | |
| def _extract_rewards(episode_data: dict[str, Any]) -> list[float]: | |
| rewards = episode_data.get("rewards") | |
| if rewards is None: | |
| rewards = episode_data.get("per_step_rewards") | |
| if rewards is None: | |
| rewards = [] | |
| return [float(reward) for reward in rewards] | |
| def _compute_score(task_id: str, rewards: list[float]) -> float: | |
| max_reward = _MAX_REWARDS.get(task_id, _MAX_REWARDS["task_1"]) | |
| total_reward = sum(rewards) | |
| raw_score = total_reward / max_reward if max_reward > 0 else 0.0 | |
| return float(max(_SCORE_MIN, min(_SCORE_MAX, raw_score))) | |
| def _make_result(task_id: str, score: float, rewards: list[float]) -> dict[str, Any]: | |
| raw_total_reward = round(sum(rewards), 4) | |
| return { | |
| "score": round(score, 4), | |
| "reward": round(score, 4), | |
| "normalized_reward": round(score, 4), | |
| "success": score >= 0.5, | |
| "task_id": task_id, | |
| "total_reward": raw_total_reward, | |
| "raw_total_reward": raw_total_reward, | |
| "steps": len(rewards), | |
| "details": { | |
| "per_step_rewards": [round(reward, 4) for reward in rewards], | |
| "max_possible_reward": _MAX_REWARDS.get(task_id, _MAX_REWARDS["task_1"]), | |
| "score_range": [_SCORE_MIN, _SCORE_MAX], | |
| }, | |
| } | |
| class BaseEpisodeRubric(OpenEnvRubric, metaclass=_RubricMeta): | |
| task_id: str = "task_1" | |
| def __init__(self) -> None: | |
| self.reset() | |
| def reset(self) -> None: | |
| self._rewards: list[float] = [] | |
| def forward(self, action: Any, observation: Any) -> float: | |
| reward_value = None | |
| if isinstance(observation, dict): | |
| reward_value = observation.get("reward") | |
| elif observation is not None and hasattr(observation, "reward"): | |
| reward_value = getattr(observation, "reward") | |
| if reward_value is not None: | |
| self._rewards.append(float(reward_value)) | |
| return _compute_score(self.task_id, self._rewards) | |
| def __call__(self, episode_data: dict[str, Any]) -> dict[str, Any]: | |
| rewards = _extract_rewards(episode_data) | |
| self._rewards = list(rewards) | |
| score = _compute_score(self.task_id, rewards) | |
| return _make_result(self.task_id, score, rewards) | |
| def grade(self, episode_data: dict[str, Any]) -> dict[str, Any]: | |
| return self(episode_data) | |
| def score(self, episode_data: dict[str, Any]) -> float: | |
| return float(self(episode_data)["score"]) | |
| class Task1Rubric(BaseEpisodeRubric): | |
| task_id = "task_1" | |
| class Task2Rubric(BaseEpisodeRubric): | |
| task_id = "task_2" | |
| class Task3Rubric(BaseEpisodeRubric): | |
| task_id = "task_3" | |
| Task1Grader = Task1Rubric | |
| Task2Grader = Task2Rubric | |
| Task3Grader = Task3Rubric | |
| GRADERS: dict[str, type[BaseEpisodeRubric]] = { | |
| "task_1": Task1Rubric, | |
| "task_2": Task2Rubric, | |
| "task_3": Task3Rubric, | |
| } | |
| def get_grader(task_id: str) -> BaseEpisodeRubric: | |
| grader_cls = GRADERS.get(task_id, Task1Rubric) | |
| return grader_cls() | |