from __future__ import annotations import math import random import uuid from typing import Any, Dict, List, Optional, Tuple from .models import ( DeliveryWindow, ManufacturingAction, ManufacturingObservation, ManufacturingReward, PendingOrder, PlatformState, ) # ─── Tuneable constants ──────────────────────────────────────────────────────── PRODUCE_MATERIAL_COST = 15.0 PRODUCE_ENERGY_COST = 10.0 PRODUCE_COMPONENT_GAIN = 12.0 ASSEMBLE_COMPONENT_COST = 10.0 ASSEMBLE_ENERGY_COST = 12.0 DELIVER_ENERGY_COST = 8.0 RECHARGE_BASE = 20.0 ENERGY_MAX = 100.0 MATERIAL_MAX = 100.0 COMPONENT_MAX = 100.0 PRODUCT_MAX = 10 MATERIAL_REFILL_RATE = 5.0 # materials added each step passively PRODUCE_MATERIAL_THRESHOLD = PRODUCE_MATERIAL_COST ASSEMBLE_COMPONENT_THRESHOLD = ASSEMBLE_COMPONENT_COST # ─── Product types ───────────────────────────────────────────────────────────── SIMPLE_PRODUCTS = ["alloy_panel", "circuit_board", "fuel_cell"] ASSEMBLED_PRODUCTS = ["solar_array", "thruster_module", "habitat_unit"] ALL_PRODUCTS = SIMPLE_PRODUCTS + ASSEMBLED_PRODUCTS class ManufacturingFactoryEnv: """Core simulation engine for the Space Manufacturing RL environment.""" def __init__( self, num_platforms: int, max_steps: int, seed: int, pending_orders: List[PendingOrder], delivery_windows: List[DeliveryWindow], solar_zones: Optional[Dict[str, float]] = None, ) -> None: self.num_platforms = num_platforms self.max_steps = max_steps self.seed = seed self._initial_orders = list(pending_orders) self._initial_windows = list(delivery_windows) self._solar_zones: Dict[str, float] = solar_zones or {"zone_a": 0.8} # mutable state (initialised in reset) self.platforms: List[PlatformState] = [] self.pending_orders: List[PendingOrder] = [] self.delivery_windows: List[DeliveryWindow] = [] self.step_count: int = 0 self.done: bool = False self.total_reward: float = 0.0 self.episode_id: str = "" self.metrics: Dict[str, Any] = {} self._rng = random.Random(seed) # ── Public API ───────────────────────────────────────────────────────────── def reset(self) -> ManufacturingObservation: self._rng = random.Random(self.seed) self.step_count = 0 self.done = False self.total_reward = 0.0 self.episode_id = str(uuid.uuid4()) self.metrics = { "production_runs": 0, "assemblies_completed": 0, "deliveries_completed": 0, "on_time_deliveries": 0, "invalid_actions": 0, "total_steps": 0, } self.platforms = [self._init_platform(i) for i in range(self.num_platforms)] self.pending_orders = [o.model_copy() for o in self._initial_orders] self.delivery_windows = [w.model_copy() for w in self._initial_windows] return self._build_observation(reward=0.0) def step( self, action: ManufacturingAction ) -> Tuple[ManufacturingObservation, ManufacturingReward, bool, Dict[str, Any]]: if self.done: obs = self._build_observation(reward=0.0) return obs, ManufacturingReward(value=0.0), True, {} step_reward = 0.0 components: Dict[str, float] = {} for pid, platform in enumerate(self.platforms): act = action.platform_actions.get(pid, "recharge") r, c = self._apply_action(platform, act) step_reward += r for k, v in c.items(): components[k] = components.get(k, 0.0) + v # Passive material refill for p in self.platforms: p.material_stock = min(MATERIAL_MAX, p.material_stock + MATERIAL_REFILL_RATE) # Energy penalty for critically low for p in self.platforms: if p.energy < 10.0: step_reward -= 1.5 components["energy_critical"] = components.get("energy_critical", 0.0) - 1.5 # Expire overdue delivery windows (soft penalty) still_open: List[DeliveryWindow] = [] for w in self.delivery_windows: if w.deadline >= self.step_count: still_open.append(w) self.delivery_windows = still_open self.step_count += 1 self.metrics["total_steps"] = self.step_count self.total_reward += step_reward if self.step_count >= self.max_steps: self.done = True elif not self.pending_orders: self.done = True # all orders delivered — stop early reward_obj = ManufacturingReward(value=step_reward, components=components) obs = self._build_observation(reward=step_reward) return obs, reward_obj, self.done, {"metrics": dict(self.metrics)} # ── Internal helpers ─────────────────────────────────────────────────────── def _init_platform(self, pid: int) -> PlatformState: angle = (2 * math.pi / self.num_platforms) * pid pos = [round(math.cos(angle) * 400, 2), round(math.sin(angle) * 400, 2), 0.0] return PlatformState( id=pid, position=pos, energy=self._rng.uniform(60.0, 90.0), material_stock=self._rng.uniform(30.0, 60.0), component_stock=0.0, product_stock=0, last_action=None, ) def _apply_action( self, platform: PlatformState, action: str ) -> Tuple[float, Dict[str, float]]: reward = 0.0 components: Dict[str, float] = {} platform.last_action = action if action == "produce": if platform.material_stock < PRODUCE_MATERIAL_THRESHOLD: reward -= 2.0 components["invalid_action"] = components.get("invalid_action", 0.0) - 2.0 self.metrics["invalid_actions"] += 1 else: platform.material_stock -= PRODUCE_MATERIAL_COST platform.component_stock = min( COMPONENT_MAX, platform.component_stock + PRODUCE_COMPONENT_GAIN ) platform.energy = max(0.0, platform.energy - PRODUCE_ENERGY_COST) reward += 3.0 components["production"] = components.get("production", 0.0) + 3.0 self.metrics["production_runs"] += 1 elif action == "assemble": if platform.component_stock < ASSEMBLE_COMPONENT_THRESHOLD: reward -= 2.0 components["invalid_action"] = components.get("invalid_action", 0.0) - 2.0 self.metrics["invalid_actions"] += 1 else: platform.component_stock -= ASSEMBLE_COMPONENT_COST platform.product_stock = min(PRODUCT_MAX, platform.product_stock + 1) platform.energy = max(0.0, platform.energy - ASSEMBLE_ENERGY_COST) reward += 5.0 components["assembly"] = components.get("assembly", 0.0) + 5.0 self.metrics["assemblies_completed"] += 1 elif action == "deliver": if platform.product_stock < 1: reward -= 2.0 components["invalid_action"] = components.get("invalid_action", 0.0) - 2.0 self.metrics["invalid_actions"] += 1 else: platform.product_stock -= 1 platform.energy = max(0.0, platform.energy - DELIVER_ENERGY_COST) self.metrics["deliveries_completed"] += 1 # Check if any window is still open → on-time bonus on_time = False for w in self.delivery_windows: if w.deadline >= self.step_count: on_time = True self.metrics["on_time_deliveries"] += 1 self.delivery_windows.remove(w) break if on_time: reward += 8.0 components["on_time_delivery"] = ( components.get("on_time_delivery", 0.0) + 8.0 ) else: reward += 3.0 components["late_delivery"] = components.get("late_delivery", 0.0) + 3.0 # Consume a pending order if matched if self.pending_orders: self.pending_orders.pop(0) elif action == "recharge": zone = self._zone_for_platform(platform) irradiance = self._solar_zones.get(zone, 0.5) gain = RECHARGE_BASE * irradiance was_low = platform.energy < 25.0 platform.energy = min(ENERGY_MAX, platform.energy + gain) if was_low: reward += 1.0 components["timely_recharge"] = components.get("timely_recharge", 0.0) + 1.0 else: # Unknown action treated as invalid reward -= 2.0 components["invalid_action"] = components.get("invalid_action", 0.0) - 2.0 self.metrics["invalid_actions"] += 1 # Idle penalty — recharge when actionable work exists if action == "recharge" and platform.energy > 50.0 and ( platform.material_stock >= PRODUCE_MATERIAL_THRESHOLD or platform.component_stock >= ASSEMBLE_COMPONENT_THRESHOLD or platform.product_stock > 0 ): reward -= 0.3 components["idle_penalty"] = components.get("idle_penalty", 0.0) - 0.3 return reward, components def _zone_for_platform(self, platform: PlatformState) -> str: zones = list(self._solar_zones.keys()) idx = platform.id % len(zones) return zones[idx] def _build_observation(self, reward: float) -> ManufacturingObservation: return ManufacturingObservation( platforms=[p.model_copy() for p in self.platforms], time_step=self.step_count, delivery_windows=list(self.delivery_windows), solar_conditions=dict(self._solar_zones), pending_orders=list(self.pending_orders), total_reward=self.total_reward, reward=reward, done=self.done, metadata=dict(self.metrics), )