Spaces:
Sleeping
Sleeping
File size: 5,509 Bytes
22805ee 12f3696 22805ee | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 | """Canonical local environment API for the satellite submission."""
from typing import Any, Dict, Tuple
from uuid import uuid4
try:
from .constellation import SatelliteConstellationEnv
from .models import (
SatelliteAction,
SatelliteEnvState,
SatelliteObservation,
SatelliteReward,
SatelliteState,
)
from .tasks import EasyTask, HardTask, MediumTask, Task
except ImportError: # pragma: no cover - flat-module fallback
from constellation import SatelliteConstellationEnv
from models import (
SatelliteAction,
SatelliteEnvState,
SatelliteObservation,
SatelliteReward,
SatelliteState,
)
from tasks import EasyTask, HardTask, MediumTask, Task
TASK_PRESETS = {
"easy": EasyTask,
"medium": MediumTask,
"hard": HardTask,
}
class SatelliteTaskEnv:
"""Canonical typed environment with reset/step/state methods."""
def __init__(
self,
task_name: str = "medium",
num_satellites: int | None = None,
max_steps: int | None = None,
):
if task_name not in TASK_PRESETS:
raise ValueError(
f"Unsupported task_name '{task_name}'. Expected one of {sorted(TASK_PRESETS)}."
)
self.task_name = task_name
self.task: Task = TASK_PRESETS[task_name]()
default_satellites = num_satellites or 5
default_max_steps = max_steps or 100
self._sim = SatelliteConstellationEnv(
num_satellites=default_satellites,
max_steps=default_max_steps,
)
self.episode_id = str(uuid4())
self.done = False
self.last_info: Dict[str, Any] = {}
self.last_reward = SatelliteReward(value=0.0, components={})
self.task.setup_environment(self._sim)
@staticmethod
def list_tasks() -> Dict[str, str]:
"""List selectable task presets."""
return {name: task_cls().description for name, task_cls in TASK_PRESETS.items()}
def set_task(self, task_name: str) -> None:
"""Switch to a different task preset."""
if task_name not in TASK_PRESETS:
raise ValueError(
f"Unsupported task_name '{task_name}'. Expected one of {sorted(TASK_PRESETS)}."
)
self.task_name = task_name
self.task = TASK_PRESETS[task_name]()
self.task.setup_environment(self._sim)
def reset(self) -> SatelliteObservation:
"""Reset the active task and return the initial observation."""
self.episode_id = str(uuid4())
self.done = False
self.last_info = {"task_name": self.task_name, "reset": True}
self.last_reward = SatelliteReward(value=0.0, components={})
self.task.setup_environment(self._sim)
observation_dict = self._sim.reset()
return self._observation_from_dict(
observation_dict,
reward=self.last_reward.value,
done=self.done,
metadata=self.last_info,
)
def step(
self, action: SatelliteAction
) -> Tuple[SatelliteObservation, SatelliteReward, bool, Dict[str, Any]]:
"""Execute one action and return (observation, reward, done, info)."""
observation_dict, reward_value, done, info = self._sim.step(action.satellite_actions)
self.done = done
self.last_info = {"task_name": self.task_name, **info}
self.last_reward = SatelliteReward(
value=reward_value,
components=info.get("reward_components", {}),
)
observation = self._observation_from_dict(
observation_dict,
reward=reward_value,
done=done,
metadata=self.last_info,
)
return observation, self.last_reward, done, self.last_info
def state(self) -> SatelliteEnvState:
"""Return a typed state snapshot for the current episode."""
observation_dict = self._sim._get_observation()
satellites = [SatelliteState(**sat) for sat in observation_dict["satellites"]]
return SatelliteEnvState(
episode_id=self.episode_id,
task_name=self.task_name,
step_count=self._sim.current_step,
max_steps=self._sim.max_steps,
seed=self._sim.seed + self._sim.episode_index,
done=self.done,
total_reward=observation_dict["total_reward"],
metrics={key: float(value) for key, value in self._sim.metrics.items()},
satellites=satellites,
ground_stations=observation_dict["ground_stations"],
weather_conditions=observation_dict["weather_conditions"],
pending_tasks=observation_dict["pending_tasks"],
)
def _observation_from_dict(
self,
observation_dict: Dict[str, Any],
reward: float,
done: bool,
metadata: Dict[str, Any],
) -> SatelliteObservation:
satellites = [SatelliteState(**sat) for sat in observation_dict["satellites"]]
return SatelliteObservation(
satellites=satellites,
time_step=observation_dict["time_step"],
ground_stations=observation_dict["ground_stations"],
weather_conditions=observation_dict["weather_conditions"],
pending_tasks=observation_dict["pending_tasks"],
total_reward=observation_dict["total_reward"],
reward=reward,
done=done,
metadata=metadata,
)
|