"""Deterministic satellite constellation simulator.""" import math from typing import Any, Dict, List, Optional, Tuple import numpy as np VALID_ACTIONS = {"capture", "downlink", "maintain", "idle"} class SatelliteConstellationEnv: """Environment for managing a satellite constellation.""" def __init__(self, num_satellites: int = 5, max_steps: int = 100, seed: int = 7): self.num_satellites = num_satellites self.max_steps = max_steps self.seed = seed self.episode_index = 0 self.rng = np.random.default_rng(seed) self.current_step = 0 self.satellites: List[Dict[str, Any]] = [] self.ground_stations = [(0, 0), (45, 90), (-30, 120)] self.weather = {"region1": 0.2, "region2": 0.5} self.capture_regions = { "region1": (18.5, 73.9), "region2": (34.0, -117.0), "region3": (-22.8, -43.2), } self.pending_tasks: List[Dict[str, Any]] = [] self.total_reward = 0.0 self.metrics: Dict[str, Any] = {} self.action_trace: List[Dict[str, Any]] = [] self.last_action_events: List[Dict[str, Any]] = [] self.no_progress_steps = 0 self._reset_metrics() self._reset_satellites() def _reset_rng(self) -> None: self.rng = np.random.default_rng(self.seed + self.episode_index) def _reset_metrics(self) -> None: self.metrics = { "successful_captures": 0, "capture_task_completions": 0, "downlink_units": 0.0, "downlink_task_completions": 0, "invalid_actions": 0, "idle_steps": 0, "maintain_actions": 0, "repeated_action_penalties": 0, "destructive_action_penalties": 0, "tasks_completed": 0, } self.action_trace = [] def _reset_satellites(self) -> None: self.satellites = [] for i in range(self.num_satellites): altitude_km = float(self.rng.uniform(420.0, 620.0)) orbit_radius = 6371.0 + altitude_km phase = float(self.rng.uniform(0.0, 2.0 * np.pi)) inclination = float(self.rng.uniform(-1.05, 1.05)) ascending_node = float(self.rng.uniform(0.0, 2.0 * np.pi)) angular_velocity = float(self.rng.uniform(0.008, 0.016)) satellite = { "id": i, "position": (0.0, 0.0, 0.0), "battery": 100.0, "storage": 0.0, "last_action": "idle", "repeat_count": 0, "orbit_radius": orbit_radius, "orbit_phase": phase, "orbit_inclination": inclination, "orbit_ascending_node": ascending_node, "orbit_angular_velocity": angular_velocity, } self._update_satellite_position(satellite) self.satellites.append(satellite) def _update_satellite_position(self, sat: Dict[str, Any]) -> None: """Project orbital parameters into 3D Cartesian space.""" radius = float(sat["orbit_radius"]) phase = float(sat["orbit_phase"]) inclination = float(sat["orbit_inclination"]) ascending_node = float(sat["orbit_ascending_node"]) cos_phase = np.cos(phase) sin_phase = np.sin(phase) cos_inc = np.cos(inclination) sin_inc = np.sin(inclination) cos_node = np.cos(ascending_node) sin_node = np.sin(ascending_node) x = radius * ((cos_node * cos_phase) - (sin_node * sin_phase * cos_inc)) y = radius * ((sin_node * cos_phase) + (cos_node * sin_phase * cos_inc)) z = radius * (sin_phase * sin_inc) sat["position"] = (float(x), float(y), float(z)) def reset(self) -> Dict[str, Any]: self.current_step = 0 self.total_reward = 0.0 self.no_progress_steps = 0 self.episode_index += 1 self._reset_rng() self._reset_metrics() self._reset_satellites() if not self.pending_tasks: self.pending_tasks = [ {"id": "img-1", "type": "image_capture", "region": "region1", "priority": 1}, { "id": "down-1", "type": "data_downlink", "station": 0, "priority": 2, "units_remaining": 20, }, ] else: self.pending_tasks = [self._clone_task(task) for task in self.pending_tasks] return self._get_observation() def step( self, action: Dict[int, str] ) -> Tuple[Dict[str, Any], float, bool, Dict[str, Any]]: self.current_step += 1 reward_value = 0.0 reward_components: Dict[str, float] = {} self.last_action_events = [] prev_tasks_completed = int(self.metrics.get("tasks_completed", 0)) prev_downlink_units = float(self.metrics.get("downlink_units", 0.0)) prev_invalid_actions = int(self.metrics.get("invalid_actions", 0)) prev_avg_battery = self._mean_battery() prev_avg_storage = self._mean_storage() for sat in self.satellites: sat_id = sat["id"] act = action.get(sat_id, "idle") if act not in VALID_ACTIONS: act = "idle" reward_value -= 1.0 self.metrics["invalid_actions"] += 1 reward_components[f"invalid_{sat_id}"] = reward_components.get(f"invalid_{sat_id}", 0.0) - 1.0 sat_reward, sat_components, action_event = self._apply_action(sat, act) reward_value += sat_reward for name, value in sat_components.items(): reward_components[name] = reward_components.get(name, 0.0) + value if action_event is not None: self.last_action_events.append(action_event) self._advance_positions() reward_value += self._apply_passive_dynamics(reward_components) reward_value += self._apply_progress_shaping( reward_components, prev_tasks_completed=prev_tasks_completed, prev_downlink_units=prev_downlink_units, prev_invalid_actions=prev_invalid_actions, prev_avg_battery=prev_avg_battery, prev_avg_storage=prev_avg_storage, ) # small stochasticity to break perfectly-constant per-step rewards noise = float(self.rng.uniform(-0.06, 0.12) + (0.02 * math.sin(self.current_step * 0.45))) reward_value += noise if abs(noise) > 1e-9: reward_components["stochastic_noise"] = reward_components.get("stochastic_noise", 0.0) + noise self.total_reward += reward_value done = self.current_step >= self.max_steps or all(s["battery"] <= 0 for s in self.satellites) observation = self._get_observation() info = { "reward_components": reward_components, "metrics": dict(self.metrics), "tasks_remaining": len(self.pending_tasks), "seed": self.seed + self.episode_index, "step": self.current_step, "action_events": list(self.last_action_events), } self.action_trace.append( { "step": self.current_step, "action": dict(action), "reward": reward_value, "reward_components": dict(reward_components), "metrics": dict(self.metrics), } ) return observation, reward_value, done, info def _apply_action( self, sat: Dict[str, Any], action: str ) -> Tuple[float, Dict[str, float], Optional[Dict[str, Any]]]: reward = 0.0 components: Dict[str, float] = {} action_event: Optional[Dict[str, Any]] = None sat_id = sat["id"] if action == sat["last_action"]: sat["repeat_count"] += 1 else: sat["repeat_count"] = 0 # only penalize very long non-idle streaks to avoid over-shaping. if sat["repeat_count"] >= 5 and action != "idle": repeat_pen = -0.15 reward += repeat_pen self.metrics["repeated_action_penalties"] += 1 components[f"repeat_penalty_{sat_id}"] = repeat_pen if action == "capture": capture_reward, action_event = self._handle_capture(sat) reward += capture_reward components[f"capture_{sat_id}"] = capture_reward elif action == "downlink": downlink_reward, action_event = self._handle_downlink(sat) reward += downlink_reward components[f"downlink_{sat_id}"] = downlink_reward elif action == "maintain": maintain_reward = self._handle_maintain(sat) reward += maintain_reward components[f"maintain_{sat_id}"] = maintain_reward else: idle_reward = self._handle_idle(sat) reward += idle_reward components[f"idle_{sat_id}"] = idle_reward # small incentive only when the chosen action is currently useful. if action != "idle" and self._is_action_useful(sat, action): proactive_bonus = 0.1 reward += proactive_bonus components[f"proactive_{sat_id}"] = components.get(f"proactive_{sat_id}", 0.0) + proactive_bonus sat["last_action"] = action return reward, {k: v for k, v in components.items() if abs(v) > 1e-9}, action_event def _handle_capture(self, sat: Dict[str, Any]) -> Tuple[float, Optional[Dict[str, Any]]]: sat_id = sat["id"] if sat["battery"] <= 12 or sat["storage"] >= 90: self.metrics["invalid_actions"] += 1 self.metrics["destructive_action_penalties"] += 1 return -1.0, None task = self._select_capture_task(sat) if task is None: if self._next_task("image_capture") is not None: miss_penalty = -0.02 - min(0.2, 0.015 * sat["repeat_count"] + 0.015 * self.no_progress_steps) return miss_penalty, None return -0.02, None cloud_cover = float(self.weather.get(task["region"], 0.5)) task_bonus = max(0.5, 3.0 * (1.0 - cloud_cover)) priority_bonus = float(task.get("priority", 1)) sat["battery"] = max(0.0, sat["battery"] - 5.0) sat["storage"] = min(100.0, sat["storage"] + 10.0) self.metrics["successful_captures"] += 1 self.metrics["capture_task_completions"] += 1 self._complete_task(task["id"]) region_name = str(task.get("region", "")) region_lat, region_lon = self.capture_regions.get(region_name, (0.0, 0.0)) return ( 4.5 + task_bonus + priority_bonus, { "satellite_id": sat_id, "action": "capture", "task_id": task["id"], "region": region_name, "target_latitude": float(region_lat), "target_longitude": float(region_lon), "units": 10.0, }, ) def _handle_downlink(self, sat: Dict[str, Any]) -> Tuple[float, Optional[Dict[str, Any]]]: sat_id = sat["id"] if sat["battery"] <= 5 or sat["storage"] <= 0: self.metrics["invalid_actions"] += 1 return -1.5, None task = self._select_downlink_task(sat) if task is None: if self._next_task("data_downlink") is not None: miss_penalty = -0.03 - min(0.25, 0.015 * sat["repeat_count"] + 0.02 * self.no_progress_steps) return miss_penalty, None return -0.03, None units_remaining = float(task.get("units_remaining", 20.0)) sent = min(sat["storage"], 20.0, units_remaining) if sent <= 0: self.metrics["invalid_actions"] += 1 return -0.5, None sat["storage"] -= sent sat["battery"] = max(0.0, sat["battery"] - 2.0) task["units_remaining"] = max(0.0, units_remaining - sent) self.metrics["downlink_units"] += sent station_id = int(task.get("station", 0)) sat_lat, sat_lon, _ = self._satellite_geo(sat) gs_lat, gs_lon = self._station_coords(station_id) distance = self._great_circle_distance_rad(sat_lat, sat_lon, gs_lat, gs_lon) distance_bonus = max(0.0, 1.0 - (distance / math.pi)) reward = sent * 1.25 + float(task.get("priority", 1)) + distance_bonus if task["units_remaining"] <= 0: self.metrics["downlink_task_completions"] += 1 self._complete_task(task["id"]) reward += 3.0 return ( reward, { "satellite_id": sat_id, "action": "downlink", "task_id": task["id"], "station_id": station_id, "target_latitude": float(gs_lat), "target_longitude": float(gs_lon), "units": float(sent), }, ) def _handle_maintain(self, sat: Dict[str, Any]) -> float: self.metrics["maintain_actions"] += 1 battery_before = sat["battery"] sat["battery"] = min(100.0, sat["battery"] + 18.0) if battery_before < 35.0: return 3.0 if battery_before < 60.0: return 1.0 self.metrics["destructive_action_penalties"] += 1 return -1.0 def _handle_idle(self, sat: Dict[str, Any]) -> float: self.metrics["idle_steps"] += 1 if self._has_actionable_visible_task(sat): return -0.05 - min(0.12, 0.015 * self.no_progress_steps) if sat["battery"] < 20: sat["battery"] = min(100.0, sat["battery"] + 1.0) return 0.15 return 0.0 def _advance_positions(self) -> None: for sat in self.satellites: sat["orbit_phase"] = float(sat["orbit_phase"] + sat["orbit_angular_velocity"]) self._update_satellite_position(sat) def _apply_passive_dynamics(self, reward_components: Dict[str, float]) -> float: battery_penalty = 0.0 storage_penalty = 0.0 for sat in self.satellites: # passive drain sat["battery"] = max(0.0, sat["battery"] - 0.5) # softer, proportional penalty for low battery if sat["battery"] < 10.0: battery_penalty -= min(1.0, (10.0 - sat["battery"]) / 20.0) # small penalty for near-full storage, scaled if sat["storage"] > 95.0: storage_penalty -= min(0.5, (sat["storage"] - 95.0) / 50.0) total_penalty = battery_penalty + storage_penalty if abs(total_penalty) > 1e-9: reward_components["resource_risk"] = reward_components.get("resource_risk", 0.0) + total_penalty return total_penalty def _apply_progress_shaping( self, reward_components: Dict[str, float], prev_tasks_completed: int, prev_downlink_units: float, prev_invalid_actions: int, prev_avg_battery: float, prev_avg_storage: float, ) -> float: task_delta = int(self.metrics.get("tasks_completed", 0)) - prev_tasks_completed downlink_delta = float(self.metrics.get("downlink_units", 0.0)) - prev_downlink_units invalid_delta = int(self.metrics.get("invalid_actions", 0)) - prev_invalid_actions avg_battery = self._mean_battery() avg_storage = self._mean_storage() battery_delta = avg_battery - prev_avg_battery storage_relief = prev_avg_storage - avg_storage shaping = 0.0 if task_delta > 0: shaping += float(task_delta) * 1.4 if downlink_delta > 0: shaping += min(2.0, 0.1 * downlink_delta) if invalid_delta > 0: shaping -= float(invalid_delta) * 0.6 if prev_avg_battery < 45.0 and battery_delta > 0: shaping += min(0.8, battery_delta / 8.0) if storage_relief > 4.0: shaping += min(0.5, storage_relief / 20.0) made_progress = (task_delta > 0) or (downlink_delta > 0.5) if made_progress: self.no_progress_steps = 0 elif self.pending_tasks: self.no_progress_steps += 1 stagnation_penalty = -min(0.5, 0.07 * self.no_progress_steps) shaping += stagnation_penalty reward_components["stagnation_penalty"] = ( reward_components.get("stagnation_penalty", 0.0) + stagnation_penalty ) else: self.no_progress_steps = 0 if abs(shaping) > 1e-9: reward_components["progress_shaping"] = reward_components.get("progress_shaping", 0.0) + shaping return shaping def _mean_battery(self) -> float: if not self.satellites: return 0.0 return float(sum(float(s["battery"]) for s in self.satellites) / len(self.satellites)) def _mean_storage(self) -> float: if not self.satellites: return 0.0 return float(sum(float(s["storage"]) for s in self.satellites) / len(self.satellites)) def _next_task(self, task_type: str) -> Optional[Dict[str, Any]]: matches = [task for task in self.pending_tasks if task["type"] == task_type] if not matches: return None matches.sort(key=lambda task: (-int(task.get("priority", 1)), str(task["id"]))) return matches[0] def _complete_task(self, task_id: str) -> None: remaining: List[Dict[str, Any]] = [] completed = False for task in self.pending_tasks: if not completed and task["id"] == task_id: completed = True continue remaining.append(task) if completed: self.pending_tasks = remaining self.metrics["tasks_completed"] += 1 def _can_downlink(self, sat_id: int) -> bool: if sat_id >= len(self.satellites): return False sat = self.satellites[sat_id] if sat["storage"] <= 0: return False return any(self._is_station_visible(sat, idx) for idx in range(len(self.ground_stations))) def _has_actionable_visible_task(self, sat: Dict[str, Any]) -> bool: can_capture = ( sat["battery"] > 12 and sat["storage"] < 90 and self._select_capture_task(sat) is not None ) can_downlink = ( sat["battery"] > 5 and sat["storage"] > 0 and self._select_downlink_task(sat) is not None ) return can_capture or can_downlink def _is_action_useful(self, sat: Dict[str, Any], action: str) -> bool: if action == "capture": return self._select_capture_task(sat) is not None if action == "downlink": return self._select_downlink_task(sat) is not None if action == "maintain": return sat["battery"] < 70 return False def _select_capture_task(self, sat: Dict[str, Any]) -> Optional[Dict[str, Any]]: candidates: List[Tuple[float, str, Dict[str, Any]]] = [] sat_lat, sat_lon, sat_alt = self._satellite_geo(sat) visibility_radius = self._visibility_radius_rad(sat_alt) for task in self.pending_tasks: if task.get("type") != "image_capture": continue region = str(task.get("region", "")) if region not in self.capture_regions: continue reg_lat, reg_lon = self.capture_regions[region] distance = self._great_circle_distance_rad(sat_lat, sat_lon, reg_lat, reg_lon) if distance > visibility_radius: continue priority = float(task.get("priority", 1)) cloud = float(self.weather.get(region, 0.5)) score = (priority * 3.0) + ((1.0 - cloud) * 2.0) - distance candidates.append((score, str(task["id"]), task)) if not candidates: return None candidates.sort(key=lambda item: (-item[0], item[1])) return candidates[0][2] def _select_downlink_task(self, sat: Dict[str, Any]) -> Optional[Dict[str, Any]]: sat_lat, sat_lon, sat_alt = self._satellite_geo(sat) visibility_radius = self._visibility_radius_rad(sat_alt) candidates: List[Tuple[float, str, Dict[str, Any]]] = [] for task in self.pending_tasks: if task.get("type") != "data_downlink": continue station_id = int(task.get("station", 0)) if station_id < 0 or station_id >= len(self.ground_stations): continue gs_lat, gs_lon = self._station_coords(station_id) distance = self._great_circle_distance_rad(sat_lat, sat_lon, gs_lat, gs_lon) if distance > visibility_radius: continue priority = float(task.get("priority", 1)) units_remaining = float(task.get("units_remaining", 20.0)) completion_bias = 0.75 if sat["storage"] >= units_remaining else 0.0 score = (priority * 3.0) + completion_bias - distance candidates.append((score, str(task["id"]), task)) if not candidates: return None candidates.sort(key=lambda item: (-item[0], item[1])) return candidates[0][2] def _station_coords(self, station_id: int) -> Tuple[float, float]: lat, lon = self.ground_stations[station_id] return float(lat), float(lon) def _is_station_visible(self, sat: Dict[str, Any], station_id: int) -> bool: sat_lat, sat_lon, sat_alt = self._satellite_geo(sat) gs_lat, gs_lon = self._station_coords(station_id) distance = self._great_circle_distance_rad(sat_lat, sat_lon, gs_lat, gs_lon) return distance <= self._visibility_radius_rad(sat_alt) def _visibility_radius_rad(self, altitude_km: float) -> float: earth_radius_km = 6371.0 alt = max(0.0, altitude_km) horizon = math.acos(min(1.0, earth_radius_km / (earth_radius_km + alt))) return max(math.radians(35.0), min(math.radians(120.0), horizon + math.radians(50.0))) def _satellite_geo(self, sat: Dict[str, Any]) -> Tuple[float, float, float]: x, y, z = sat["position"] radius = math.sqrt((x * x) + (y * y) + (z * z)) if radius <= 0: return 0.0, 0.0, 0.0 lat = math.degrees(math.asin(z / radius)) lon = math.degrees(math.atan2(y, x)) altitude = max(0.0, radius - 6371.0) return float(lat), float(lon), float(altitude) def _great_circle_distance_rad( self, lat1: float, lon1: float, lat2: float, lon2: float ) -> float: lat1_rad = math.radians(lat1) lon1_rad = math.radians(lon1) lat2_rad = math.radians(lat2) lon2_rad = math.radians(lon2) d_lat = lat2_rad - lat1_rad d_lon = lon2_rad - lon1_rad a = ( math.sin(d_lat / 2.0) ** 2 + math.cos(lat1_rad) * math.cos(lat2_rad) * math.sin(d_lon / 2.0) ** 2 ) return 2.0 * math.asin(min(1.0, math.sqrt(a))) def _clone_task(self, task: Dict[str, Any]) -> Dict[str, Any]: return {key: value for key, value in task.items()} def _get_observation(self) -> Dict[str, Any]: return { "satellites": [ { "id": s["id"], "position": s["position"], "battery": s["battery"], "storage": s["storage"], "last_action": s["last_action"], } for s in self.satellites ], "time_step": self.current_step, "ground_stations": self.ground_stations, "capture_regions": dict(self.capture_regions), "weather_conditions": self.weather, "pending_tasks": [self._clone_task(task) for task in self.pending_tasks], "total_reward": self.total_reward, }