import gymnasium as gym import numpy as np from collections import deque import tensorflow as tf from keras import optimizers from keras.optimizers import Adam import tensorflow_probability as tfp import os from datetime import datetime import json from RaceCar import ActorCritic from Environment_Wrapper import CarRacingEnvWrapper from gymnasium.wrappers import TimeLimit tfd = tfp.distributions class PPOAgent: def __init__(self, env_id="CarRacing-v3", num_envs=21, gamma=0.99, lam=0.95, clip_epsilon=0.2, actor_lr=3e-4, critic_lr=3e-4, ppo_epochs=10, minibatches=4, steps_per_batch=1024, num_stack_frames=4, resize_dim=(84, 84), grayscale=True, seed=42, log_dir="./ppo_logs", entropy_coeff=0.01, save_interval_timesteps=537600, hidden_layer_sizes=[512, 512, 512]): self.env_id = env_id self.num_envs = num_envs self.gamma = gamma self.lam = lam self.clip_epsilon = clip_epsilon self.ppo_epochs = ppo_epochs self.minibatches = minibatches self.steps_per_batch = steps_per_batch self.num_stack_frames = num_stack_frames self.resize_dim = resize_dim self.grayscale = grayscale self.seed = seed self.log_dir = log_dir self.entropy_coeff = entropy_coeff self.save_interval_timesteps = save_interval_timesteps self.hidden_layer_sizes = hidden_layer_sizes self.envs = self._make_vec_envs() self.action_dim = self.envs.single_action_space.shape[0] self.observation_shape = self.envs.single_observation_space.shape self.model = ActorCritic(self.action_dim, self.num_stack_frames, self.resize_dim[1], self.resize_dim[0], hidden_layer_sizes=self.hidden_layer_sizes) self.actor_optimizer = Adam(learning_rate=actor_lr) self.critic_optimizer = Adam(learning_rate=critic_lr) self.train_log_dir = None self.summary_writer = None tf.random.set_seed(self.seed) np.random.seed(self.seed) dummy_input = np.zeros((1, *self.observation_shape), dtype=np.uint8) _ = self.model(dummy_input) def _make_env(self): env = gym.make(self.env_id, render_mode="rgb_array", continuous=True) env = CarRacingEnvWrapper(env, num_stack_frames=self.num_stack_frames, grayscale=self.grayscale, resize_dim=self.resize_dim) env = TimeLimit(env, max_episode_steps=1000) return env def _make_vec_envs(self): return gym.vector.AsyncVectorEnv([lambda: self._make_env() for _ in range(self.num_envs)]) def _compute_returns_and_advantages(self, rewards, values, dones): advantages = np.zeros_like(rewards, dtype=np.float32) returns = np.zeros_like(rewards, dtype=np.float32) last_gae_lam = np.zeros(self.num_envs, dtype=np.float32) for t in reversed(range(self.steps_per_batch)): next_non_terminal = 1.0 - dones[t] delta = rewards[t] + self.gamma * values[t+1] * next_non_terminal - values[t] last_gae_lam = delta + self.gamma * self.lam * next_non_terminal * last_gae_lam advantages[t] = last_gae_lam returns = advantages + values[:-1] return advantages, returns @tf.function def _train_step(self, observations, actions, old_log_probs, advantages, returns): with tf.GradientTape() as tape: action_distribution, value_pred = self.model(observations) value_pred = tf.squeeze(value_pred, axis=-1) critic_loss = tf.reduce_mean(tf.square(returns - value_pred)) log_prob = action_distribution.log_prob(actions) ratio = tf.exp(log_prob - old_log_probs) ratio = tf.where(tf.math.is_nan(ratio), 1.0, ratio) ratio = tf.where(tf.math.is_inf(ratio), tf.sign(ratio) * 1e5, ratio) pg_loss1 = ratio * advantages pg_loss2 = tf.clip_by_value(ratio, 1 - self.clip_epsilon, 1 + self.clip_epsilon) * advantages actor_loss = -tf.reduce_mean(tf.minimum(pg_loss1, pg_loss2)) entropy = tf.reduce_mean(action_distribution.entropy()) entropy_loss = -self.entropy_coeff * entropy total_loss = actor_loss + critic_loss + entropy_loss grads = tape.gradient(total_loss, self.model.trainable_variables) self.actor_optimizer.apply_gradients(zip(grads, self.model.trainable_variables)) return actor_loss, critic_loss, entropy, total_loss def train(self, total_timesteps, resume_from_timestep=0, resume_model_path=None, run_log_dir=None): global_timestep = resume_from_timestep ep_rewards = deque(maxlen=100) ep_modified_rewards = deque(maxlen=100) resume_json_path = "resume_config.json" if run_log_dir is None: if os.path.exists(resume_json_path): with open(resume_json_path, "r") as f: resume_info = json.load(f) self.train_log_dir = resume_info.get("run_log_directory") if self.train_log_dir is None: current_time = datetime.now().strftime("%Y%m%d-%H%M%S") self.train_log_dir = os.path.join(self.log_dir, current_time) else: self.train_log_dir = run_log_dir os.makedirs(os.path.join(self.train_log_dir, "checkpoints"), exist_ok=True) self.summary_writer = tf.summary.create_file_writer(self.train_log_dir) if resume_model_path and os.path.exists(resume_model_path): self.model.load_weights(resume_model_path) print(f"Resuming training from timestep {global_timestep} and loaded model from: {resume_model_path}") elif resume_from_timestep > 0: print(f"WARNING: Attempting to resume from timestep {global_timestep} but no valid model path provided or found. Starting fresh.") else: print("Starting new training run.") obs, _ = self.envs.reset(seed=self.seed) print(f"Target total timesteps: {total_timesteps}") print(f"Current global timestep: {global_timestep}") print("Starting training loop...") while global_timestep < total_timesteps: batch_observations = np.zeros((self.steps_per_batch, self.num_envs, *self.observation_shape), dtype=np.uint8) batch_actions = np.zeros((self.steps_per_batch, self.num_envs, self.action_dim), dtype=np.float32) batch_rewards = np.zeros((self.steps_per_batch, self.num_envs), dtype=np.float32) batch_dones = np.zeros((self.steps_per_batch, self.num_envs), dtype=bool) batch_values = np.zeros((self.steps_per_batch, self.num_envs), dtype=np.float32) batch_log_probs = np.zeros((self.steps_per_batch, self.num_envs), dtype=np.float32) batch_original_rewards = np.zeros((self.steps_per_batch, self.num_envs), dtype=np.float32) for i in range(self.steps_per_batch): #print("Step:", i) tf_obs = tf.convert_to_tensor(obs, dtype=tf.uint8) action_dist, value = self.model(tf_obs) #print("Model forward pass complete.") action = action_dist.sample() log_prob = action_dist.log_prob(action) #print("Action sampled and log_prob calculated.") action_np = action.numpy() value_np = tf.squeeze(value).numpy() log_prob_np = log_prob.numpy() next_obs, reward, terminated, truncated, info = self.envs.step(action_np) #print("Environment step complete.") done = terminated | truncated batch_observations[i] = obs batch_actions[i] = action_np batch_rewards[i, :] = reward batch_dones[i] = done batch_values[i] = value_np batch_log_probs[i] = log_prob_np for single_env_info in info: if isinstance(single_env_info, dict) and single_env_info.get('episode'): episode_reward = single_env_info['episode']['r'] ep_rewards.append(episode_reward) obs = next_obs global_timestep += self.num_envs _, last_values_np = self.model(tf.convert_to_tensor(obs, dtype=tf.uint8)) last_values_np = tf.squeeze(last_values_np).numpy() batch_values = np.concatenate((batch_values, last_values_np[np.newaxis, :]), axis=0) advantages, returns = self._compute_returns_and_advantages( batch_rewards, batch_values, batch_dones ) flat_observations = batch_observations.reshape((-1, *self.observation_shape)) flat_actions = batch_actions.reshape((-1, self.action_dim)) flat_old_log_probs = batch_log_probs.reshape(-1) flat_advantages = advantages.reshape(-1) flat_returns = returns.reshape(-1) flat_advantages = (flat_advantages - np.mean(flat_advantages)) / (np.std(flat_advantages) + 1e-8) batch_original_total_reward = np.sum(batch_original_rewards) batch_modified_total_reward = np.sum(batch_rewards) batch_indices = np.arange(self.steps_per_batch * self.num_envs) for _ in range(self.ppo_epochs): np.random.shuffle(batch_indices) for start_idx in range(0, len(batch_indices) // self.minibatches * self.minibatches, len(batch_indices) // self.minibatches): end_idx = start_idx + len(batch_indices) // self.minibatches minibatch_indices = batch_indices[start_idx:end_idx] mb_obs = tf.constant(flat_observations[minibatch_indices], dtype=tf.uint8) mb_actions = tf.constant(flat_actions[minibatch_indices], dtype=tf.float32) mb_old_log_probs = tf.constant(flat_old_log_probs[minibatch_indices], dtype=tf.float32) mb_advantages = tf.constant(flat_advantages[minibatch_indices], dtype=tf.float32) mb_returns = tf.constant(flat_returns[minibatch_indices], dtype=tf.float32) actor_loss, critic_loss, entropy, total_loss = self._train_step( mb_obs, mb_actions, mb_old_log_probs, mb_advantages, mb_returns ) with self.summary_writer.as_default(): tf.summary.scalar("charts/total_timesteps", global_timestep, step=global_timestep) tf.summary.scalar("losses/actor_loss", actor_loss, step=global_timestep) tf.summary.scalar("losses/critic_loss", critic_loss, step=global_timestep) tf.summary.scalar("losses/entropy", entropy, step=global_timestep) tf.summary.scalar("losses/total_loss", total_loss, step=global_timestep) if ep_rewards: tf.summary.scalar("charts/avg_episode_reward_original_env", np.mean(ep_rewards), step=global_timestep) tf.summary.scalar("charts/min_episode_reward_original_env", np.min(ep_rewards), step=global_timestep) tf.summary.scalar("charts/max_episode_reward_original_env", np.max(ep_rewards), step=global_timestep) tf.summary.scalar("charts/batch_total_reward_original", batch_original_total_reward, step=global_timestep) tf.summary.scalar("charts/batch_total_reward_modified", batch_modified_total_reward, step=global_timestep) with tf.name_scope('actor_std'): tf.summary.histogram('log_std', self.model.actor_log_std, step=global_timestep) tf.summary.scalar('std_mean_overall', tf.reduce_mean(tf.exp(self.model.actor_log_std)), step=global_timestep) self.summary_writer.flush() if global_timestep % self.save_interval_timesteps == 0: if ep_rewards: avg_orig_reward_str = f"{np.mean(ep_rewards):.2f}" else: avg_orig_reward_str = 'N/A' print(f"Timestep: {global_timestep}, Number of episodes in Timestep: (525 eps)") #{avg_orig_reward_str}") current_checkpoint_path = os.path.join(self.train_log_dir, "checkpoints", f"Actor-Critic_at_{global_timestep}.weights.h5") self.model.save_weights(current_checkpoint_path) resume_info = { "last_global_timestep": global_timestep, "last_checkpoint_path": current_checkpoint_path, "run_log_directory": self.train_log_dir } with open(resume_json_path, "w") as f: json.dump(resume_info, f, indent=4) print(f"Resume info saved to {resume_json_path}") print("Training finished.") self.envs.close() self.summary_writer.close() final_model_path = os.path.join(self.train_log_dir, "Actor-Critic_final_model.weights.h5") self.model.save_weights(final_model_path) if os.path.exists(resume_json_path): os.remove(resume_json_path) print(f"Removed {resume_json_path} as training completed successfully.") def evaluate(self, num_episodes=5, render=True, model_path=None): eval_env = gym.make(self.env_id, render_mode="human" if render else "rgb_array", continuous=True) eval_env = CarRacingEnvWrapper(eval_env, num_stack_frames=self.num_stack_frames, grayscale=self.grayscale, resize_dim=self.resize_dim) if model_path: dummy_input = np.zeros((1, *self.observation_shape), dtype=np.uint8) _ = self.model(dummy_input) self.model.load_weights(model_path) print(f"Loaded model from {model_path}") episode_rewards = [] episode_original_rewards = [] for ep in range(num_episodes): obs, _ = eval_env.reset() done = False total_reward = 0 total_original_reward = 0 while not done: tf_obs = tf.convert_to_tensor(obs[np.newaxis, :], dtype=tf.uint8) action_dist, _ = self.model(tf_obs) action = action_dist.mean().numpy().flatten() action[0] = np.clip(action[0], -1.0, 1.0) action[1] = np.clip(action[1], 0.0, 1.0) action[2] = np.clip(action[2], 0.0, 1.0) obs, reward, terminated, truncated, info = eval_env.step(action) done = terminated or truncated total_reward += reward total_original_reward += info.get('original_reward', reward) episode_rewards.append(total_reward) episode_original_rewards.append(total_original_reward) print(f"Episode {ep+1} finished. Modified Reward: {total_reward:.2f}, Original Env Reward: {total_original_reward:.2f}") eval_env.close() print(f"Average modified evaluation reward over {num_episodes} episodes: {np.mean(episode_rewards):.2f}") print(f"Average original environment reward over {num_episodes} episodes: {np.mean(episode_original_rewards):.2f}") return episode_rewards, episode_original_rewards