Spaces:
Sleeping
Sleeping
| from __future__ import annotations | |
| from functools import lru_cache | |
| from .environment import V2SupplyMindEnv | |
| from .models import V2TaskResult | |
| from .policies import heuristic_joint_policy, naive_joint_policy | |
| from .solver import RolloutStats, rollout_reference_stats | |
| STRICT_SCORE_EPSILON = 1e-4 | |
| BASELINE_SCORE_ANCHOR = 0.05 | |
| REFERENCE_SCORE_ANCHOR = 0.95 | |
| def grade_episode(task_id: str, seed: int, raw_reward: float, center_reward: float, average_warehouse_reward: float) -> V2TaskResult: | |
| baseline = cached_rollout_policy(task_id, seed, "baseline") | |
| reference = cached_reference_stats(task_id, seed) | |
| target_reward = max(reference.global_reward, baseline.global_reward + 20.0) | |
| target_center_reward = max(reference.center_reward, baseline.center_reward + 5.0) | |
| target_warehouse_reward = max(reference.average_warehouse_reward, baseline.average_warehouse_reward + 5.0) | |
| return V2TaskResult( | |
| task_id=task_id, | |
| raw_reward=raw_reward, | |
| baseline_reward=baseline.global_reward, | |
| target_reward=target_reward, | |
| score=normalize_score(raw_reward, baseline.global_reward, target_reward), | |
| center_reward=center_reward, | |
| average_warehouse_reward=average_warehouse_reward, | |
| baseline_center_reward=baseline.center_reward, | |
| target_center_reward=target_center_reward, | |
| center_role_score=normalize_score(center_reward, baseline.center_reward, target_center_reward), | |
| baseline_warehouse_reward=baseline.average_warehouse_reward, | |
| target_warehouse_reward=target_warehouse_reward, | |
| warehouse_role_score=normalize_score(average_warehouse_reward, baseline.average_warehouse_reward, target_warehouse_reward), | |
| ) | |
| def rollout_policy(task_id: str, seed: int, policy_name: str) -> RolloutStats: | |
| env = V2SupplyMindEnv(default_task_id=task_id) | |
| observation = env.reset_internal(task_id, seed) | |
| policy = naive_joint_policy if policy_name == "baseline" else heuristic_joint_policy | |
| while not env.done: | |
| result = env.step(policy(observation), grade_terminal=False) | |
| observation = result.observation | |
| warehouse_rewards = [value for key, value in env.agent_rewards.items() if key != "center"] | |
| return RolloutStats( | |
| global_reward=env.cumulative_reward, | |
| center_reward=env.agent_rewards["center"], | |
| average_warehouse_reward=sum(warehouse_rewards) / len(warehouse_rewards) if warehouse_rewards else 0.0, | |
| ) | |
| def cached_rollout_policy(task_id: str, seed: int, policy_name: str) -> RolloutStats: | |
| return rollout_policy(task_id, seed, policy_name) | |
| def cached_reference_stats(task_id: str, seed: int) -> RolloutStats: | |
| return rollout_reference_stats(task_id, seed) | |
| def normalize_score(raw_reward: float, baseline_reward: float, target_reward: float) -> float: | |
| lower = STRICT_SCORE_EPSILON | |
| upper = 1.0 - STRICT_SCORE_EPSILON | |
| if target_reward <= baseline_reward: | |
| return REFERENCE_SCORE_ANCHOR if raw_reward >= target_reward else BASELINE_SCORE_ANCHOR | |
| progress = (raw_reward - baseline_reward) / (target_reward - baseline_reward) | |
| if progress <= 1: | |
| score = BASELINE_SCORE_ANCHOR + progress * (REFERENCE_SCORE_ANCHOR - BASELINE_SCORE_ANCHOR) | |
| else: | |
| score = REFERENCE_SCORE_ANCHOR + min(progress - 1, 1.0) * (upper - REFERENCE_SCORE_ANCHOR) | |
| return max(lower, min(upper, score)) | |