ZacBl/OrbitRL
0
1import pygame2import math3import random4import csv5 6import torch 7 8from collections import deque9 10from Agent import *11from Class_Def import *12from utils import *13from Const import *14 15device = DEVICE16 17pygame.init()18 19WIN = pygame.display.set_mode((WIDTH, HEIGHT))20pygame.display.set_caption("Planet_Sim")21FONT = pygame.font.SysFont("Times", 16)22 23def SatSim(load_weights=False):24 print(device)25 Scale = 50/AU26 Step_p_frame = 127 shift_x, shift_y = 0, 028 run = True29 clock = pygame.time.Clock()30 31 planets = Solar_System()32 Target = 6 #random.randint(0,len(planets)-1)33 rockets = []34 rockets.append(Launch(planets[3]))35 36 37 # --- Initialize Distance Recording ---38 all_distance_acquisitions = []39 current_acquisition = []40 41 # --- Instantiate the Agent ---42 agent = Att_Agent(n_features=NUM_FEATURES, # Use n_features43 seq_len=SEQ_LENGTH, # Use seq_len44 action_size=ACTION_SIZE,45 replay_memory_capacity=MEMORY_CAPACITY,46 batch_size=BATCH_SIZE,47 gamma=GAMMA,48 eps_start=EPS_START, eps_end=EPS_END, eps_decay=EPS_DECAY,49 tau=TAU, lr=LR, target_update_freq=TARGET_UPDATE_FREQ,50 # Add Attention Network specific params if needed (using defaults here)51 # embed_dim=64, num_heads=452 )53 54 # Load weights if demanded55 try:56 agent.load_weights()57 except Exception as e:58 print(f"Error loading weights: {e}")59 # --- Setup Agent Logging ---60 # The agent will create its own timestamped directory within Metrics/MODEL_NAME/61 agent.setup_logging(model_name=MODEL_NAME) # Use the defined MODEL_NAME62 63 # --- Initialize Episode Tracking Variables ---64 episode_number = 065 episode_reward = 0.066 episode_steps = 067 episode_goal_achieved = False68 69 # History deques initialization...70 action_hist = deque(maxlen=MAX_HISTORY_LEN)71 72 Flag = 073 agent_controlled_rocket_index = 0 74 rocket = rockets[agent_controlled_rocket_index] if rockets else None75 target_planet = planets[Target] if Target < len(planets) else None76 77 78 # --- Initialize History Deques ---79 action_hist = deque(maxlen=MAX_HISTORY_LEN)80 rel_x_hist = deque(maxlen=MAX_HISTORY_LEN)81 rel_y_hist = deque(maxlen=MAX_HISTORY_LEN)82 rel_vx_hist = deque(maxlen=MAX_HISTORY_LEN)83 rel_vy_hist = deque(maxlen=MAX_HISTORY_LEN)84 85 # Fill initial history with padding values (optional, but helps)86 # Use padding values defined in utils or Agent87 initial_pad_value = 0.088 initial_action_pad = 089 for _ in range(SEQ_LENGTH): # Start with enough history for the first state90 action_hist.append(initial_action_pad)91 rel_x_hist.append(initial_pad_value)92 rel_y_hist.append(initial_pad_value)93 rel_vx_hist.append(initial_pad_value)94 rel_vy_hist.append(initial_pad_value)95 96 # Initialize state tuple for the loop start97 state_tuple = None98 99 # --- Steps count ---100 n_step = 0 101 while run:102 # --- Initialize state ---103 # For simplicity, let's always control the first rocket if it exists104 done = False # Reset done flag each iteration105 106 agent_controlled_rocket_index = 0 107 rocket = rockets[agent_controlled_rocket_index] if rockets else None108 target_planet = planets[Target]109 110 clock.tick(FPS)111 WIN.fill((0, 0, 0))112 113 114 # --- Pygame Event Handling ---115 for event in pygame.event.get():116 if event.type == pygame.QUIT:117 run = False118 #Actions done once per pressure119 if event.type == pygame.KEYDOWN:120 if(event.key == pygame.K_1): Step_p_frame += 10 # Speed up time fast >>121 if(event.key == pygame.K_2): Step_p_frame += 1 # Speed up time slow >122 if(event.key == pygame.K_3): Step_p_frame -= 1 # Slow down time slow <123 if(event.key == pygame.K_4): Step_p_frame -= 10 # Slow down time fast <<124 if(event.key == pygame.K_a): Flag -=1 # Show previous object125 if(event.key == pygame.K_z): Flag +=1 # Show next object126 if(event.key == pygame.K_t): Flag = Target # Show target object127 if(event.key == pygame.K_r): Target = random.randint(0,len(planets)-1) # Set a new target128 if(event.key == pygame.K_SPACE):129 rockets.append(Launch(planets[3]))130 131 if(event.key == pygame.K_BACKSPACE):132 if(len(rockets)> 0):133 rockets.pop()134 if event.key == pygame.K_x: 135 if rocket:136 done = True137 138 # rocket = reset_rocket_state(rocket, planets) Now reset is in the done condition True139 # --- Clear History Deques on Reset ---140 action_hist.clear()141 rel_x_hist.clear()142 rel_y_hist.clear()143 rel_vx_hist.clear()144 rel_vy_hist.clear()145 # Fill with padding again? Or let the loop handle it?146 # Re-filling ensures immediate valid sequence length.147 for _ in range(SEQ_LENGTH):148 action_hist.append(initial_action_pad)149 rel_x_hist.append(initial_pad_value)150 rel_y_hist.append(initial_pad_value)151 rel_vx_hist.append(initial_pad_value)152 rel_vy_hist.append(initial_pad_value)153 # state = get_state(rocket, target_planet) # Remove old state logic154 state_tuple = None # Indicate state needs recalculation155 # --- Log previous episode before reset (if it ran at least one step) ---156 # if episode_steps > 0:157 # agent.log_episode_data(episode_number, episode_steps, episode_reward, episode_goal_achieved)158 # --- Reset Episode Trackers ---159 episode_reward = 0.0160 episode_steps = 0161 episode_goal_achieved = False162 print("Targeted planet is", Core_position[int(Target)])163 # Save distance track and reset164 # all_distance_acquisitions.append(list(current_acquisition)) # Use list() to copy165 # current_acquisition.clear() # Use clear() for deque/list166 167 if(event.key == pygame.K_q): 168 run = False169 all_distance_acquisitions.append(current_acquisition) # Save the track of distances bwt Rckt & Tgt170 current_acquisition = []171 172 if(Flag>(len(planets)+len(rockets)-1)): Flag = 0173 if(Flag<0): Flag = len(planets)+len(rockets)-1174 175 # --- RL Agent Step ---176 reward = 0 # Initialize reward177 178 if rocket and target_planet: # Ensure we have a rocket and target179 # --- 1. Calculate Current Relative State ---180 # This provides the LATEST snapshot before action selection181 # Uses the state *before* physics update for this step182 current_rel_x, current_rel_y, current_rel_vx, current_rel_vy = calculate_current_relative_state(rocket, target_planet)183 184 # Handle potential NaN from calculate_current_relative_state if rocket/target missing185 if any(math.isnan(v) for v in [current_rel_x, current_rel_y, current_rel_vx, current_rel_vy]):186 print("Warning: Invalid current state values (NaN). Skipping agent step.")187 # Potentially try to recover or skip agent logic for this frame188 # You might need error handling here depending on how NaNs occur189 else:190 # --- 2. Get State Sequence for Agent ---191 # This uses the history *leading up to* the current point192 state_tuple = get_state_sequence(193 action_hist, rel_x_hist, rel_y_hist, rel_vx_hist, rel_vy_hist, seq_length=SEQ_LENGTH194 )195 196 # --- 3. Agent Selects Action ---197 action_tensor = agent.select_action(state_tuple) # Pass tuple (sequence, mask)198 action = action_tensor.item()199 200 # --- 4. Apply Action to Rocket ---201 rocket.apply_action(action)202 203 # --- Store state before physics update (for reward calc based on distance change) ---204 prev_dist = get_distance(rocket, target_planet)205 if prev_dist == 0: prev_dist = 1e-6 # Avoid division by zero206 207 # --- 5. Physics Update ---208 planets, rockets = update_position(planets, rockets) # This updates rocket.x, .y, .vx, .vy209 210 # Find the controlled rocket again after update (it might have been removed)211 # This assumes the agent always controls the rocket at index 0 if it exists212 controlled_rocket_index = 0213 if controlled_rocket_index < len(rockets):214 rocket = rockets[controlled_rocket_index] # Update rocket reference215 else:216 rocket = None # Rocket was destroyed/removed217 print("Controlled rocket destroyed during physics update.")218 done = True219 reward = -GOAL_REWARD * 2 # Heavy penalty for destruction220 next_state_tuple = None # No next state if destroyed221 222 223 # --- 6. Calculate Reward & Next State (if rocket exists) ---224 if rocket:225 # --- Increment episode step counter ---226 episode_steps += 1227 try:228 current_distance = get_distance(rocket, target_planet)229 current_acquisition.append(current_distance) # Save distance for logging230 231 if current_distance == 0: current_distance = 1e-6232 233 # --- 6a. Calculate Reward Components ---234 # Distance Reward :235 reward_distance = float(AU/abs(current_distance - 10*target_planet.radius)) 236 if current_distance>50*target_planet.radius: 237 reward_distance *= np.sign(prev_dist - current_distance)238 239 240 # Orbital speed reward & goal check241 target_mass = target_planet.mass242 ideal_orbital_speed = math.sqrt(G * target_mass / max(current_distance, 1e-6)) # Avoid div by zero243 rel_vx_now = rocket.vx - target_planet.vx # Use updated velocities244 rel_vy_now = rocket.vy - target_planet.vy245 current_speed_relative = math.sqrt(rel_vx_now**2 + rel_vy_now**2)246 speed_diff = abs((current_speed_relative - ideal_orbital_speed) / current_speed_relative)247 reward_speed = REWARD_SPEED_SCALE * (1.0 - speed_diff)248 249 reward_action = 0250 if current_distance<20*target_planet.radius:251 reward_action = NO_THRUST_REWARD if action == 0 else 0 # Penalize any thrust252 253 # Time Penalty254 reward_time = 0 #-TIME_PENALTY (Far too soon) Probably useless w/ fuel consuption notion.255 256 # Goal Reward Check (use updated distance/speed)257 is_orbiting_goal = (speed_diff * ideal_orbital_speed < ORBIT_SPEED_TOLERANCE) and \258 (5*target_planet.radius <= current_distance <= 20*target_planet.radius)259 # print(current_distance, target_planet.radius)260 reactors_off = rocket.motor_off()261 262 if is_orbiting_goal and reactors_off:263 episode_goal_achieved = True # Mark true if goal met anytime in episode264 reward_goal = GOAL_REWARD265 else:266 reward_goal = 0267 268 reward = reward_distance + reward_speed + reward_action + reward_time + reward_goal269 step_reward = reward270 271 # --- Accumulate Episode Reward ---272 episode_reward += step_reward273 274 # --- 6b. Check "Done" Conditions ---275 # Crash into target planet276 crash_dist_m = (target_planet.radius) #+ CRASH_DISTANCE_THRESHOLD277 if current_distance < crash_dist_m :278 print("Crashed into target! Resetting rocket.")279 print(current_distance, target_planet.radius)280 done = True281 reward -= GOAL_REWARD * 0.5 # Penalty for crash282 # Resetting logic now handles history clearing (see point 5)283 284 285 # Out of bounds (relative to target planet)286 if current_distance > OUT_OF_BOUNDS_DISTANCE:287 print("Went out of bounds! Resetting rocket.")288 done = True289 reward -= GOAL_REWARD # Penalty for OOB (adjust as needed)290 # Resetting logic handles history clearing291 292 293 # --- 6c. Update History & Get Next State ---294 # Calculate the relative state AFTER the physics update295 next_rel_x, next_rel_y, next_rel_vx, next_rel_vy = calculate_current_relative_state(rocket, target_planet)296 297 # Append the state values LEADING TO the action, and the action itself, to history298 # Use the values calculated in step 1299 if not math.isnan(current_rel_x): # Check if state was valid before appending300 action_hist.append(action)301 rel_x_hist.append(current_rel_x)302 rel_y_hist.append(current_rel_y)303 rel_vx_hist.append(current_rel_vx)304 rel_vy_hist.append(current_rel_vy)305 306 # Now generate the next_state sequence using the updated history307 if not done and not any(math.isnan(v) for v in [next_rel_x, next_rel_y, next_rel_vx, next_rel_vy]):308 next_state_tuple = get_state_sequence(309 action_hist, rel_x_hist, rel_y_hist, rel_vx_hist, rel_vy_hist, seq_length=SEQ_LENGTH310 )311 else:312 next_state_tuple = None # Terminal state or invalid next state313 314 315 except (NameError, AttributeError, IndexError, TypeError, ValueError) as e:316 print(f"Warning: Error during RL step calculation: {e}")317 import traceback318 traceback.print_exc() # Print full traceback for debugging319 reward = 0 # Default reward on error320 next_state_tuple = None # Cannot determine next state321 # Decide if this error should terminate the episode (done=True)322 # done = True # Optional: Terminate on calculation error323 324 325 # --- 7. Store Experience ---326 # Ensure state_tuple was calculated in this iteration before storing327 if state_tuple is not None:328 reward_tensor = torch.tensor([reward], dtype=torch.float32, device=device)329 # Store the state tuple (history before action) and next_state tuple (history after action)330 agent.store_experience(state_tuple, action_tensor, next_state_tuple, reward_tensor, done)331 else:332 # This case might happen on the very first step if state isn't pre-initialized333 # or after an error/reset where state_tuple became None.334 # Avoid storing experience if the initial state wasn't valid.335 print("Warning: state_tuple is None, skipping experience storage.") # Optional log336 pass337 338 339 # --- 8. Optimize Agent Model ---340 agent.optimize_model()341 342 # --- 9. Update Target Network ---343 if agent.tau > 0:344 agent.update_target_net(soft_update=True)345 # --- Or Hard update less frequently ---346 # elif step_count % agent.target_update_freq == 0:347 # agent.update_target_net(soft_update=False)348 349 350 # --- Handle End of Episode ---351 if done:352 print(f"Episode finished. Reason: {'Crash' if current_distance < crash_dist_m else 'OOB' if current_distance > OUT_OF_BOUNDS_DISTANCE else 'User Reset'}. Final Reward: {reward}")353 # Clear History Deques on Reset354 action_hist.clear(); rel_x_hist.clear(); rel_y_hist.clear(); rel_vx_hist.clear(); rel_vy_hist.clear()355 # Re-fill padding356 for _ in range(SEQ_LENGTH):357 action_hist.append(initial_action_pad)358 rel_x_hist.append(initial_pad_value)359 rel_y_hist.append(initial_pad_value)360 rel_vx_hist.append(initial_pad_value)361 rel_vy_hist.append(initial_pad_value)362 # Save distance track and reset363 all_distance_acquisitions.append(list(current_acquisition))364 current_acquisition.clear()365 episode_number += 1366 episode_reward /= episode_steps # Average reward per step367 print(f"--- Episode {episode_number} Finished --- Steps: {episode_steps}, Average Reward: {episode_reward:.2f}, Goal: {episode_goal_achieved} ---")368 # Log data for the completed episode369 agent.log_episode_data(episode_number, episode_steps, episode_reward, episode_goal_achieved)370 371 # Reset rocket state372 rocket = reset_rocket_state(rocket, planets)373 # Clear and refill history deques374 action_hist.clear(); rel_x_hist.clear(); rel_y_hist.clear(); rel_vx_hist.clear(); rel_vy_hist.clear()375 for _ in range(SEQ_LENGTH):376 action_hist.append(initial_action_pad); rel_x_hist.append(initial_pad_value); rel_y_hist.append(initial_pad_value); rel_vx_hist.append(initial_pad_value); rel_vy_hist.append(initial_pad_value)377 378 # Reset episode trackers379 episode_reward = 0.0380 episode_steps = 0381 episode_goal_achieved = False382 state_tuple = None # Invalidate state tuple until next loop iteration383 # --- End Reward Calculation ---384 385 386 387 388 # Add a pannel to show commands when pressing 'h':389 keys = pygame.key.get_pressed()390 if keys[pygame.K_o]: Scale *= 1.05 # Zoom in391 if keys[pygame.K_p]: Scale /= 1.05 # Zoom out392 393 if keys[pygame.K_h]:394 help_lines = [395 "--- Help ---",396 "1: Speed up time (fast)",397 "2: Speed up time (slow)",398 "3: Slow down time (slow)",399 "4: Slow down time (fast)",400 "a: Focus previous object",401 "z: Focus next object",402 "t: Focus target planet",403 "r: Set random target planet",404 "space: Launch new rocket (from Earth)",405 "backspace: Remove last rocket",406 "o: Zoom in",407 "p: Zoom out",408 "left/right/up: Activate rocket reactors (when focusing rocket)"409 ]410 411 line_height = FONT.get_height()412 start_y = HEIGHT - (len(help_lines) * line_height) - 30 # Start Y position near bottom413 414 for i, line in enumerate(help_lines):415 line_surface = FONT.render(line, 1, WHITE)416 # Position each line below the previous one417 WIN.blit(line_surface, (10, start_y + i * line_height))418 419 420 421 422 speed_text = FONT.render(f"Step per Frame: {int(Step_p_frame)}", 1, WHITE)423 WIN.blit(speed_text, (WIDTH - speed_text.get_width() - 10, 25))424 425 follow_core = Core_position[int(Flag)] if int(Flag)<len(planets) else f"Rocket {int(Flag-len(planets))}"426 427 Flag_text = FONT.render(f"Followed Planet: {follow_core}", 1, WHITE)428 WIN.blit(Flag_text, (10, 10))429 430 Target_text = FONT.render(f"Target Planet: {Core_position[int(Target)]}", 1, WHITE)431 WIN.blit(Target_text, (10, 25))432 433 reward_text_line1 = f"Total Reward: {round(reward, 3)}"434 reward_text_line2 = f" D: {round(reward_distance, 3)}, S: {round(reward_speed, 3)}, A: {round(reward_action, 3)}, T: {round(reward_time, 3)}, G: {round(reward_goal, 3)}"435 reward_text1 = FONT.render(reward_text_line1, 1, WHITE)436 reward_text2 = FONT.render(reward_text_line2, 1, WHITE)437 WIN.blit(reward_text1, (10, 40))438 WIN.blit(reward_text2, (10, 55)) # Position below the first line439 440 Dist_text = FONT.render(f"Rocket-Tgt Distancce:{round(get_distance(planets[Target],rockets[0])/AU,5)}", 1, WHITE)441 WIN.blit(Dist_text, (WIDTH-Dist_text.get_width()-10, HEIGHT-Dist_text.get_height()-25)) 442 443 if(Flag<len(planets)):444 Orbit_text = FONT.render(f"Is Orbiting:{is_orbiting(planets[3],planets[Flag])}", 1, WHITE)445 else:446 Orbit_text = FONT.render(f"Is Orbiting:{is_orbiting(rockets[Flag-len(planets)],planets[Target])}", 1, WHITE)447 WIN.blit(Orbit_text, (WIDTH-Orbit_text.get_width()-10, HEIGHT-Dist_text.get_height()-10))448 # in the bottom right corner : show press h to show help449 help_text = FONT.render("Press h to show help-commands", 1, WHITE)450 WIN.blit(help_text, (10, HEIGHT - help_text.get_height() - 10))451 # --- Show fps ---452 fps = clock.get_fps()453 fps_text = FONT.render(f"FPS: {int(fps)}", 1, WHITE)454 WIN.blit(fps_text, (WIDTH - fps_text.get_width() - 10, HEIGHT-Dist_text.get_height()-40))455 456 # --- Update Display ---457 if n_step%Step_p_frame == 0:458 # --- Draw Planets and Rockets ---459 if Flag<len(planets):460 shift_x = - planets[Flag].x461 shift_y = - planets[Flag].y462 else:463 shift_x = - rockets[Flag-len(planets)].x464 shift_y = - rockets[Flag-len(planets)].y465 for core in planets+rockets:466 core.draw(WIN, shift_x, shift_y, Scale)467 pygame.display.update()468 n_step += 1469 pygame.quit()470 print("Simulation ended.")471 # Optional: Save the trained model472 agent.save_weights_and_distances(all_distance_acquisitions)473 # Close the log file474 agent.close_log()475 476 477if __name__ == '__main__':478 479 480 SatSim(load_weights=True)