| 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 |
|
|
|
|
| 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() |
| |
| dummy_state = tf.zeros((1, state_dim), dtype=tf.float32) |
| self.actor(dummy_state) |
| self.critic(dummy_state) |
|
|
| |
| 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 |
| ) |
|
|
| |
| 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 |
|
|
| |
| @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) |
|
|
| |
| @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) |
| 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)) |
| 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), |
| } |
