import os import tensorflow as tf from model import Actor, Critic, gaussian_log_prob, gaussian_entropy from config import ( LEARNING_RATE, LR_MIN_FRACTION, GAMMA, GAE_LAMBDA, CLIP_RATIO, VALUE_LOSS_COEF, INITIAL_ENTROPY_COEF, MAX_GRAD_NORM, CHECKPOINT_DIR, TOTAL_ITERATIONS ) from utils import compute_gae_vectorized def _build_optimizer(optimizer, variables): """Create optimizer slots without taking a real update step.""" try: optimizer.build(variables) except AttributeError: pass # legacy optimizers create slots lazily; checkpoint restore is deferred instead def _set_lr(optimizer, value): try: optimizer.learning_rate.assign(value) except Exception: optimizer.learning_rate = value class PPOAgent: def __init__(self, state_dim, action_dim): self.gamma = GAMMA self.gae_lambda = GAE_LAMBDA self.base_clip_ratio = CLIP_RATIO self.value_loss_coef = VALUE_LOSS_COEF self.entropy_coef = INITIAL_ENTROPY_COEF self.global_step = tf.Variable(0, trainable=False, dtype=tf.int64, name="global_step") self.actor = Actor(action_dim) self.critic = Critic() # Build the networks so every variable (incl. log_std) exists before the optimizers dummy_state = tf.zeros((1, state_dim), dtype=tf.float32) self.actor(dummy_state) self.critic(dummy_state) # Fail loudly if log_std is not trained (sigma would silently stay frozen) names = [v.name for v in self.actor.trainable_variables] if not any("log_std" in n for n in names): raise RuntimeError(f"actor log_std is not trainable! variables: {names}") print(f"Actor trainable variables: {len(names)} (log_std tracked)") self.actor_optimizer = tf.keras.optimizers.Adam(learning_rate=LEARNING_RATE) self.critic_optimizer = tf.keras.optimizers.Adam(learning_rate=LEARNING_RATE) _build_optimizer(self.actor_optimizer, self.actor.trainable_variables) _build_optimizer(self.critic_optimizer, self.critic.trainable_variables) self.checkpoint = tf.train.Checkpoint( global_step=self.global_step, actor=self.actor, critic=self.critic, actor_optimizer=self.actor_optimizer, critic_optimizer=self.critic_optimizer, ) self.checkpoint_manager = tf.train.CheckpointManager( self.checkpoint, directory=CHECKPOINT_DIR, max_to_keep=5 ) # schedules / checkpoints def update_schedules(self, current_iteration, total_iterations=TOTAL_ITERATIONS): """Track the global step and linearly anneal the learning rate.""" self.global_step.assign(current_iteration) frac = max(1.0 - (current_iteration - 1) / float(total_iterations), LR_MIN_FRACTION) self.current_lr = LEARNING_RATE * frac _set_lr(self.actor_optimizer, self.current_lr) _set_lr(self.critic_optimizer, self.current_lr) @property def current_clip_ratio(self): return self.base_clip_ratio def save_checkpoint(self): path = self.checkpoint_manager.save() print(f"Checkpoint saved at step {self.global_step.numpy()}: {path}") def load_checkpoint(self): latest = self.checkpoint_manager.latest_checkpoint if latest: self.checkpoint.restore(latest) print(f"Restored from checkpoint (resuming at global step {self.global_step.numpy()}): {latest}") return True print("No checkpoint found. Starting fresh from global step 0.") return False def save_era(self, directory): """Write a self-contained checkpoint into its own folder (no file copying).""" os.makedirs(directory, exist_ok=True) return self.checkpoint.save(os.path.join(directory, "ckpt")) def load_from(self, directory): """Restore the newest checkpoint inside `directory` (used for eras / inference).""" path = tf.train.latest_checkpoint(directory) if path is None: return False self.checkpoint.restore(path).expect_partial() return True # acting @tf.function(reduce_retracing=True) def act_batch(self, states): """Sample actions for all envs at once. Returns the RAW Gaussian sample (stored for training), its log-prob and V(s). Clip to [-1, 1] only when stepping the env.""" mu, sigma = self.actor(states) raw = mu + sigma * tf.random.normal(tf.shape(mu)) log_prob = gaussian_log_prob(raw, mu, sigma) value = tf.squeeze(self.critic(states), axis=-1) return raw, log_prob, value @tf.function(reduce_retracing=True) def value_batch(self, states): return tf.squeeze(self.critic(states), axis=-1) def compute_trajectory_advantages(self, rewards, values, dones, next_values): return compute_gae_vectorized(rewards, values, dones, next_values, self.gamma, self.gae_lambda) # learning @tf.function(reduce_retracing=True) def train_step(self, states, actions, old_log_probs, returns, advantages, entropy_coef, clip_val): with tf.GradientTape() as actor_tape: mu, sigma = self.actor(states) log_probs = gaussian_log_prob(actions, mu, sigma) log_ratio = log_probs - old_log_probs ratio = tf.exp(log_ratio) surr1 = ratio * advantages surr2 = tf.clip_by_value(ratio, 1.0 - clip_val, 1.0 + clip_val) * advantages policy_loss = -tf.reduce_mean(tf.minimum(surr1, surr2)) entropy = gaussian_entropy(sigma) # scalar: summed over action dims actor_loss = policy_loss - entropy_coef * entropy actor_grads = actor_tape.gradient(actor_loss, self.actor.trainable_variables) actor_grads, _ = tf.clip_by_global_norm(actor_grads, MAX_GRAD_NORM) self.actor_optimizer.apply_gradients(zip(actor_grads, self.actor.trainable_variables)) with tf.GradientTape() as critic_tape: values_pred = tf.squeeze(self.critic(states), axis=-1) value_loss = tf.reduce_mean(tf.square(returns - values_pred)) # no value clipping critic_loss = self.value_loss_coef * value_loss critic_grads = critic_tape.gradient(critic_loss, self.critic.trainable_variables) critic_grads, _ = tf.clip_by_global_norm(critic_grads, MAX_GRAD_NORM) self.critic_optimizer.apply_gradients(zip(critic_grads, self.critic.trainable_variables)) return { "policy_loss": policy_loss, "value_loss": value_loss, "entropy": entropy, "approx_kl": tf.reduce_mean((ratio - 1.0) - log_ratio), "clip_frac": tf.reduce_mean(tf.cast(tf.abs(ratio - 1.0) > clip_val, tf.float32)), "sigma_mean": tf.reduce_mean(sigma), }