| import gymnasium as gym |
| from gymnasium.vector import AsyncVectorEnv |
| import numpy as np |
| import tensorflow as tf |
| import json |
| import os |
| import time |
|
|
| from reward_shaping import LunarLanderRewardShaping |
| from config import * |
|
|
| def save_resume_data(filepath, timesteps, episodes): |
| """ |
| Saves the current training state to a JSON file located at the specified filepath. |
| """ |
| data = { |
| "timesteps": timesteps, |
| "episode_count": episodes, |
| "date": time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()) |
| } |
| try: |
| with open(filepath, 'w') as f: |
| json.dump(data, f) |
| except Exception as e: |
| print(f"Error saving resume data to {filepath}: {e}") |
|
|
| def load_resume_data(filepath): |
| """ |
| Loads the last saved training state from the specified JSON file. |
| Returns (timesteps, episode_count). |
| """ |
| if os.path.exists(filepath): |
| try: |
| with open(filepath, 'r') as f: |
| data = json.load(f) |
| timesteps = data.get("timesteps", 0) |
| episodes = data.get("episode_count", 0) |
| return (timesteps, episodes) |
| except Exception as e: |
| print(f"Error loading resume data from {filepath}: {e}. Starting from scratch.") |
| return (0, 0) |
|
|
| def make_env(env_id, seed, idx, **kwargs): |
| """ |
| Creates a single environment instance with a unique seed and applies necessary wrappers. |
| """ |
| def thunk(): |
| env = gym.make(env_id, **kwargs) |
| |
| env = LunarLanderRewardShaping(env) |
| |
| env = gym.wrappers.RecordEpisodeStatistics(env) |
| |
| env.action_space.seed(seed + idx) |
| env.observation_space.seed(seed + idx) |
| |
| return env |
| return thunk |
|
|
| def make_parallel_envs(env_id, num_envs, seed): |
| """ |
| Creates multiple environments and wraps them in an AsyncVectorEnv. |
| """ |
| env_fns = [make_env(env_id, seed, i) for i in range(num_envs)] |
| return AsyncVectorEnv(env_fns) |
|
|
| def calculate_gae(rewards, values, terminated, truncated, next_value, gamma=GAMMA, gae_lambda=GAE_LAMBDA): |
| """ |
| Calculates Generalized Advantage Estimation (GAE) and Returns (R) from rollout data. |
| """ |
| advantages = np.zeros_like(rewards, dtype=np.float32) |
| last_gae_lambda = 0 |
| |
| for t in reversed(range(N_STEPS)): |
| done_mask = 1.0 - (terminated[t] | truncated[t]).astype(np.float32) |
| if t == N_STEPS - 1: |
| next_non_terminal_value = next_value |
| next_value_actual = values[t] |
| else: |
| next_non_terminal_value = values[t + 1] |
| next_value_actual = values[t] |
| |
| delta = rewards[t] + gamma * next_non_terminal_value * done_mask - next_value_actual |
| |
| advantages[t] = delta + gamma * gae_lambda * done_mask * last_gae_lambda |
| last_gae_lambda = advantages[t] |
| |
| returns = advantages + values |
| return advantages, returns |
