SAT / constellation.py
KUMARISHWETA242
Final Commit
03fa3e4
Raw
History Blame Contribute Delete
24.1 kB
"""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,
}