| import csv |
| import json |
| import os |
| import time |
| import numpy as np |
| import numba |
| from numba import njit |
| import gymnasium as gym |
| import psutil |
| from collections import deque |
|
|
| @njit |
| def ppo_epoch_jit(H_re, H_im, W_re, W_im, Wc_re, Wc_im, a_vec, adv_n, ret_, old_lps, log_std, log_std_grad_out, bias, perm, scale, sqrtD, actor_lr, critic_lr, clip_eps, bias_lr, ent_coef, action_dim): |
| T = H_re.shape[0] |
| A = action_dim |
| D = H_re.shape[1] |
| LOG_2PI = numba.float32(1.8378770664093453) |
| policy_loss_sum = numba.float64(0.0) |
| value_loss_sum = numba.float64(0.0) |
| for ii in range(T): |
| idx = perm[ii] |
| Hr = H_re[idx] |
| Hi = H_im[idx] |
| mu = np.empty(A, numba.float32) |
| for k in range(A): |
| s = numba.float32(0.0) |
| for d in range(D): |
| s += W_re[k, d] * Hr[d] + W_im[k, d] * Hi[d] |
| mu[k] = s / scale |
| new_lp = numba.float32(0.0) |
| for k in range(A): |
| sigma = np.exp(log_std[k]) |
| z = (a_vec[idx, k] - mu[k]) / sigma |
| new_lp += numba.float32(-0.5) * (z * z + numba.float32(2.0) * log_std[k] + LOG_2PI) |
| ratio = np.exp(new_lp - old_lps[idx]) |
| if ratio > numba.float32(20.0): |
| ratio = numba.float32(20.0) |
| adv = adv_n[idx] |
| lo = numba.float32(1.0) - clip_eps |
| hi2 = numba.float32(1.0) + clip_eps |
| rc = ratio if ratio < hi2 else hi2 |
| rc = rc if rc > lo else lo |
| surr_u = ratio * adv |
| surr_c = rc * adv |
| surr = surr_u if surr_u < surr_c else surr_c |
| policy_loss_sum += numba.float64(-surr) |
| COEF_CLIP = numba.float32(2.0) |
| Z_CLIP = numba.float32(5.0) |
| for k in range(A): |
| sigma = np.exp(log_std[k]) |
| diff = a_vec[idx, k] - mu[k] |
| mean_score = diff / (sigma * sigma) |
| coef = actor_lr * surr * mean_score / scale |
| if coef > COEF_CLIP: |
| coef = COEF_CLIP |
| elif coef < -COEF_CLIP: |
| coef = -COEF_CLIP |
| for d in range(D): |
| W_re[k, d] += coef * Hr[d] |
| W_im[k, d] += coef * Hi[d] |
| z = diff / sigma |
| if z > Z_CLIP: |
| z = Z_CLIP |
| elif z < -Z_CLIP: |
| z = -Z_CLIP |
| log_std_grad_out[k] += surr * (z * z - numba.float32(1.0)) + ent_coef |
| target = ret_[idx] |
| bias = bias + bias_lr * (target - bias) |
| v = numba.float32(0.0) |
| for d in range(D): |
| v += Wc_re[d] * Hr[d] + Wc_im[d] * Hi[d] |
| v /= sqrtD |
| td_error = target - bias - v |
| value_loss_sum += numba.float64(td_error * td_error) |
| res = td_error * (critic_lr / sqrtD) |
| if res > COEF_CLIP: |
| res = COEF_CLIP |
| elif res < -COEF_CLIP: |
| res = -COEF_CLIP |
| for d in range(D): |
| Wc_re[d] += res * Hr[d] |
| Wc_im[d] += res * Hi[d] |
| return (bias, policy_loss_sum / T, value_loss_sum / T) |
|
|
| @njit |
| def compute_gae_jit(rewards, values, dones, last_val, gamma, lam): |
| n = rewards.shape[0] |
| adv = np.zeros(n, numba.float64) |
| gae = 0.0 |
| nv = last_val |
| for t in range(n - 1, -1, -1): |
| m = 1.0 - dones[t] |
| gae = rewards[t] + gamma * nv * m - values[t] + gamma * lam * m * gae |
| adv[t] = gae |
| nv = values[t] |
| return (adv, adv + values) |
|
|
| def warmup_jit(D, T, action_dim): |
| Hr = np.zeros((T, D), np.float32) |
| Hi = np.zeros((T, D), np.float32) |
| Wr = np.zeros((action_dim, D), np.float32) |
| Wi = np.zeros((action_dim, D), np.float32) |
| Cr = np.zeros(D, np.float32) |
| Ci = np.zeros(D, np.float32) |
| av = np.zeros((T, action_dim), np.float32) |
| an = np.zeros(T, np.float32) |
| rt = np.zeros(T, np.float32) |
| ol = np.zeros(T, np.float32) |
| ls = np.zeros(action_dim, np.float32) |
| lsg = np.zeros(action_dim, np.float32) |
| pm = np.arange(T, dtype=np.int32) |
| sc = np.float32(np.sqrt(D)) |
| sq = np.float32(np.sqrt(D)) |
| ppo_epoch_jit(Hr, Hi, Wr, Wi, Cr, Ci, av, an, rt, ol, ls, lsg, np.float32(0.0), pm, sc, sq, np.float32(0.001), np.float32(0.005), np.float32(0.2), np.float32(0.05), np.float32(0.01), action_dim) |
| compute_gae_jit(np.zeros(T, np.float64), np.zeros(T, np.float64), np.zeros(T, np.float64), 0.0, 0.99, 0.95) |
|
|
| def pendulum_features(obs): |
| cos_t, sin_t, td = (float(obs[0]), float(obs[1]), float(obs[2])) |
| return np.array([cos_t, sin_t, td, td / 8.0, td * td, sin_t * td, cos_t * td], dtype=np.float32) |
| OBS_DIM = 7 |
| ACTION_DIM = 1 |
| BASE_BETA = 2.5 |
| PENDULUM_CONFIG = dict(feat_lo=[-1.0, -1.0, -8.0, -1.0, 0.0, -8.0, -8.0], feat_hi=[1.0, 1.0, 8.0, 1.0, 64.0, 8.0, 8.0], feature_fn=pendulum_features, action_dim=ACTION_DIM, action_low=-2.0, action_high=2.0, D=512, beta=BASE_BETA, rollout_steps=1024, actor_lr=0.003, critic_lr=0.005, n_epochs=8, log_std_init=-0.5, log_std_min=-2.0, log_std_max=0.7, log_std_lr=0.01, entropy_coef=0.1, entropy_decay=1.0, entropy_min=0.1, clip_eps=0.2, gamma=0.95, lam=0.95, solve_thresh=-200.0, ema_interval=100, ema_alpha=0.15) |
| REWARD_THRESHOLDS = [-1200, -800, -500, -300, -200, -150, -120, -100] |
| DEFAULT_SEED = 123 |
| BETA_EFF_MIN_MULT = 0.2 |
| BETA_EFF_MAX_MULT = 4.0 |
| G_EMA_DECAY = 0.99 |
|
|
| class HDEncoderGradAdaptiveContinuous: |
|
|
| def __init__(self, feat_lo, feat_hi, D, seed, feature_fn, beta_base, phi_init=None): |
| self.lo = np.array(feat_lo, np.float32) |
| self.hi = np.array(feat_hi, np.float32) |
| self.feature_fn = feature_fn |
| self.beta_base = float(beta_base) |
| self.D = D |
| self.n_feat = len(feat_lo) |
| self.sqrtD = float(np.sqrt(D)) |
| if phi_init is not None: |
| self.Phi = np.asarray(phi_init, dtype=np.float32) |
| else: |
| rng = np.random.default_rng(seed) |
| self.Phi = rng.uniform(-np.pi, np.pi, (self.n_feat, D)).astype(np.float32) |
| scale_norm = 2.0 / (self.hi - self.lo + 1e-08) |
| self.dtheta_ds_unit = self.Phi * scale_norm[:, None] |
| self._actor = None |
| self._log_g_ema = 0.0 |
| self.beta_eff_history = [] |
|
|
| def link_actor(self, actor): |
| self._actor = actor |
|
|
| @property |
| def beta_vec(self): |
| return np.full(self.D, self.beta_base, dtype=np.float32) |
|
|
| def encode(self, state): |
| s = self.feature_fn(state) |
| s = np.clip(s, self.lo, self.hi) |
| s_norm = (2.0 * (s - self.lo) / (self.hi - self.lo + 1e-08) - 1.0).astype(np.float32) |
| proj = s_norm @ self.Phi |
| theta_base = self.beta_base * proj |
| H_re_base, H_im_base = (np.cos(theta_base), np.sin(theta_base)) |
| W_re, W_im = (self._actor.W_re, self._actor.W_im) |
| dtheta_ds = self.beta_base * self.dtheta_ds_unit |
| dH_re_ds = -H_im_base[None, :] * dtheta_ds |
| dH_im_ds = H_re_base[None, :] * dtheta_ds |
| J = (W_re @ dH_re_ds.T + W_im @ dH_im_ds.T) / self.sqrtD |
| g = float(np.linalg.norm(J)) |
| log_g = np.log1p(g) |
| centered = log_g - self._log_g_ema |
| self._log_g_ema = G_EMA_DECAY * self._log_g_ema + (1.0 - G_EMA_DECAY) * log_g |
| beta_eff = self.beta_base * (1.0 + centered) |
| beta_eff = float(np.clip(beta_eff, self.beta_base * BETA_EFF_MIN_MULT, self.beta_base * BETA_EFF_MAX_MULT)) |
| self.beta_eff_history.append(beta_eff) |
| theta = beta_eff * proj |
| return (np.cos(theta).astype(np.float32), np.sin(theta).astype(np.float32)) |
|
|
| class HDActorContinuous: |
|
|
| def __init__(self, D, action_dim, log_std_init, action_low, action_high): |
| self.sqrtD = float(np.sqrt(D)) |
| self.action_dim = action_dim |
| self.W_re = np.zeros((action_dim, D), dtype=np.float32) |
| self.W_im = np.zeros((action_dim, D), dtype=np.float32) |
| self.log_std = np.full(action_dim, log_std_init, dtype=np.float32) |
| self.a_lo = float(action_low) |
| self.a_hi = float(action_high) |
|
|
| def mean(self, H_re, H_im): |
| return (self.W_re @ H_re + self.W_im @ H_im) / self.sqrtD |
|
|
| def sample(self, H_re, H_im): |
| mu = self.mean(H_re, H_im) |
| sigma = np.exp(self.log_std) |
| a_raw = mu + sigma * np.random.standard_normal(self.action_dim).astype(np.float32) |
| a_env = np.clip(a_raw, self.a_lo, self.a_hi).astype(np.float32) |
| z = (a_raw - mu) / sigma |
| lp = float(np.sum(-0.5 * (z * z + 2.0 * self.log_std + np.log(2.0 * np.pi)))) |
| return (a_raw, a_env, lp) |
|
|
| def greedy(self, H_re, H_im): |
| mu = self.mean(H_re, H_im) |
| return np.clip(mu, self.a_lo, self.a_hi).astype(np.float32) |
|
|
| def snapshot(self): |
| return (self.W_re.copy(), self.W_im.copy()) |
|
|
| def restore(self, snap, alpha): |
| wr, wi = snap |
| self.W_re = (1 - alpha) * self.W_re + alpha * wr |
| self.W_im = (1 - alpha) * self.W_im + alpha * wi |
|
|
| class HDCritic: |
|
|
| def __init__(self, D, bias_lr=0.05): |
| self.sqrtD = float(np.sqrt(D)) |
| self.W_re = np.zeros(D, dtype=np.float32) |
| self.W_im = np.zeros(D, dtype=np.float32) |
| self.bias = 0.0 |
| self.bias_lr = bias_lr |
|
|
| def value(self, H_re, H_im): |
| return float(np.dot(self.W_re, H_re) + np.dot(self.W_im, H_im)) / self.sqrtD + self.bias |
|
|
| class TrajectoryBuffer: |
|
|
| def __init__(self): |
| self.reset() |
|
|
| def reset(self): |
| self.H_res, self.H_ims = ([], []) |
| self.actions, self.rewards = ([], []) |
| self.log_probs, self.values = ([], []) |
| self.dones = [] |
|
|
| def store(self, H_re, H_im, a, r, lp, v, done): |
| self.H_res.append(H_re) |
| self.H_ims.append(H_im) |
| self.actions.append(a) |
| self.rewards.append(r) |
| self.log_probs.append(lp) |
| self.values.append(v) |
| self.dones.append(done) |
|
|
| def __len__(self): |
| return len(self.rewards) |
|
|
| def to_arrays(self): |
| return (np.array(self.H_res, dtype=np.float32), np.array(self.H_ims, dtype=np.float32), np.array(self.actions, dtype=np.float32), np.array(self.rewards, dtype=np.float64), np.array(self.log_probs, dtype=np.float32), np.array(self.values, dtype=np.float64), np.array(self.dones, dtype=np.float64)) |
|
|
| class HDPPOAgentContinuous: |
|
|
| def __init__(self, cfg, seed=DEFAULT_SEED): |
| D = cfg['D'] |
| self.encoder = HDEncoderGradAdaptiveContinuous(cfg['feat_lo'], cfg['feat_hi'], D, seed, cfg['feature_fn'], cfg['beta'], phi_init=cfg.get('fpe_phi_init')) |
| self.actor = HDActorContinuous(D, cfg['action_dim'], cfg['log_std_init'], cfg['action_low'], cfg['action_high']) |
| self.encoder.link_actor(self.actor) |
| self.critic = HDCritic(D) |
| self.buffer = TrajectoryBuffer() |
| self.cfg = cfg |
| self.entropy_coef = cfg['entropy_coef'] |
| self.best_avg100 = -np.inf |
| self.best_snap = None |
|
|
| def select_action(self, state): |
| H_re, H_im = self.encoder.encode(state) |
| a_raw, a_env, lp = self.actor.sample(H_re, H_im) |
| val = self.critic.value(H_re, H_im) |
| return (a_raw, a_env, lp, val, H_re.copy(), H_im.copy()) |
|
|
| def store(self, H_re, H_im, a_raw, r, lp, v, done): |
| self.buffer.store(H_re, H_im, a_raw, r, lp, v, done) |
|
|
| def update(self, last_state, last_done): |
| buf = self.buffer |
| if not len(buf): |
| return (None, None) |
| cfg = self.cfg |
| H_re, H_im, a_vec, rewards, old_lps, values, dones = buf.to_arrays() |
| T = len(rewards) |
| if last_done: |
| last_val = 0.0 |
| else: |
| lr, li = self.encoder.encode(last_state) |
| last_val = self.critic.value(lr, li) |
| adv, returns = compute_gae_jit(rewards, values, dones, last_val, cfg['gamma'], cfg['lam']) |
| adv_std = float(adv.std()) |
| adv_n = np.clip((adv - adv.mean()) / (adv_std + 1e-08), -3.0, 3.0) if adv_std > 0.0001 else np.zeros(T, np.float64) |
| adv_n32 = adv_n.astype(np.float32) |
| ret32 = returns.astype(np.float32) |
| scale = np.float32(self.actor.sqrtD) |
| sqrtD = np.float32(self.critic.sqrtD) |
| act_lr = np.float32(cfg['actor_lr']) |
| crit_lr = np.float32(cfg['critic_lr']) |
| clip_e = np.float32(cfg['clip_eps']) |
| ent_c = np.float32(self.entropy_coef) |
| bias_lr = np.float32(self.critic.bias_lr) |
| bias = np.float32(self.critic.bias) |
| log_std = self.actor.log_std |
| policy_losses, value_losses = ([], []) |
| for _ in range(cfg['n_epochs']): |
| perm = np.random.permutation(T).astype(np.int32) |
| ls_grad = np.zeros(cfg['action_dim'], np.float32) |
| bias, pl, vl = ppo_epoch_jit(H_re, H_im, self.actor.W_re, self.actor.W_im, self.critic.W_re, self.critic.W_im, a_vec, adv_n32, ret32, old_lps, log_std, ls_grad, bias, perm, scale, sqrtD, act_lr, crit_lr, clip_e, bias_lr, ent_c, cfg['action_dim']) |
| policy_losses.append(float(pl)) |
| value_losses.append(float(vl)) |
| ls_step = cfg['log_std_lr'] * (ls_grad / np.float32(T)) |
| log_std = log_std + ls_step.astype(np.float32) |
| np.clip(log_std, cfg['log_std_min'], cfg['log_std_max'], out=log_std) |
| self.actor.log_std = log_std |
| self.critic.bias = float(bias) |
| self.buffer.reset() |
| return (float(np.mean(policy_losses)), float(np.mean(value_losses))) |
|
|
| def maybe_snapshot(self, avg100): |
| if avg100 > self.best_avg100: |
| self.best_avg100 = avg100 |
| self.best_snap = self.actor.snapshot() |
|
|
| def ema_restore(self): |
| if self.best_snap is not None: |
| self.actor.restore(self.best_snap, self.cfg['ema_alpha']) |
|
|
| class SystemMonitor: |
|
|
| def __init__(self): |
| self.proc = psutil.Process(os.getpid()) |
|
|
| def ram_mb(self): |
| return self.proc.memory_info().rss / 1024 ** 2 |
|
|
| def evaluate_agent(agent, n_episodes=20, seed_base=10000): |
| env = gym.make('Pendulum-v1') |
| rewards = np.empty(n_episodes, dtype=np.float64) |
| for i in range(n_episodes): |
| state, _ = env.reset(seed=seed_base + i) |
| ep_r = 0.0 |
| done = False |
| while not done: |
| Hr, Hi = agent.encoder.encode(state) |
| a = agent.actor.greedy(Hr, Hi) |
| state, reward, term, trunc, _ = env.step(a) |
| ep_r += float(reward) |
| done = term or trunc |
| rewards[i] = ep_r |
| env.close() |
| n = len(rewards) |
| sem = float(rewards.std(ddof=1) / np.sqrt(n)) if n > 1 else 0.0 |
| return dict(mean_reward=float(rewards.mean()), ci95_reward=1.96 * sem, n_episodes=n_episodes, rewards=rewards.tolist()) |
|
|
| def prune_actor_global(checkpoint_npz, D_prime): |
| W_re, W_im = (checkpoint_npz['W_actor_re'], checkpoint_npz['W_actor_im']) |
| Wc_re, Wc_im = (checkpoint_npz['W_critic_re'], checkpoint_npz['W_critic_im']) |
| Phi = checkpoint_npz['fpe_phi'] |
| beta_base = float(checkpoint_npz['beta_base']) |
| log_std = checkpoint_npz['log_std'] |
| importance = np.sqrt((W_re ** 2).sum(axis=0) + (W_im ** 2).sum(axis=0)) |
| keep_idx = np.sort(np.argsort(-importance)[:D_prime]) |
| return dict(D=D_prime, beta=beta_base, log_std=log_std, fpe_phi=Phi[:, keep_idx], W_actor_re=W_re[:, keep_idx], W_actor_im=W_im[:, keep_idx], W_critic_re=Wc_re[keep_idx], W_critic_im=Wc_im[keep_idx], critic_bias=float(checkpoint_npz.get('critic_bias', 0.0))) |
|
|
| def train_one_seed(seed, total_timesteps, save_weights_path=None, warm_start=None, log_csv_path=None, eval_csv_path=None, eval_every_n_steps=None, D=None, verbose=True): |
| cfg = dict(PENDULUM_CONFIG) |
| if warm_start is not None: |
| cfg['D'] = int(warm_start['D']) |
| cfg['beta'] = warm_start.get('beta', cfg['beta']) |
| cfg['fpe_phi_init'] = np.asarray(warm_start['fpe_phi'], dtype=np.float32) |
| elif D is not None: |
| cfg['D'] = int(D) |
| csv_file = csv_writer = None |
| if log_csv_path is not None: |
| csv_file = open(log_csv_path, 'w', newline='') |
| csv_writer = csv.writer(csv_file) |
| csv_writer.writerow(['global_step', 'wall_time_sec', 'episode', 'episodes_this_update', 'ep_rew_mean', 'ep_rew_max', 'ep_rew_min', 'best_avg100', 'policy_loss', 'value_loss', 'log_std_mean', 'entropy_coef', 'fps', 'ram_mb']) |
| eval_csv_file = eval_csv_writer = None |
| if eval_csv_path is not None: |
| eval_csv_file = open(eval_csv_path, 'w', newline='') |
| eval_csv_writer = csv.writer(eval_csv_file) |
| eval_csv_writer.writerow(['global_step', 'eval_mean', 'eval_ci95', 'tag']) |
| env = gym.make('Pendulum-v1') |
| np.random.seed(seed) |
| agent = HDPPOAgentContinuous(cfg, seed=seed) |
| if warm_start is not None: |
| agent.actor.W_re = np.asarray(warm_start['W_actor_re'], dtype=np.float32).copy() |
| agent.actor.W_im = np.asarray(warm_start['W_actor_im'], dtype=np.float32).copy() |
| agent.critic.W_re = np.asarray(warm_start['W_critic_re'], dtype=np.float32).copy() |
| agent.critic.W_im = np.asarray(warm_start['W_critic_im'], dtype=np.float32).copy() |
| agent.critic.bias = float(warm_start.get('critic_bias', 0.0)) |
| agent.actor.log_std = np.asarray(warm_start['log_std'], dtype=np.float32).copy() |
| if verbose: |
| print(f' Warm-started actor from provided checkpoint: D={cfg['D']}') |
| print(' Warm-started critic from provided checkpoint (warm-start, not reset)') |
| print(f' Warm-started log_std from provided checkpoint (warm-start, not reset): {agent.actor.log_std}') |
| if eval_csv_writer is not None: |
| post_prune_eval = evaluate_agent(agent) |
| eval_csv_writer.writerow([0, post_prune_eval['mean_reward'], post_prune_eval['ci95_reward'], 'post_prune']) |
| eval_csv_file.flush() |
| if verbose: |
| print(f' Post-prune eval (before fine-tuning): {post_prune_eval['mean_reward']:+.1f} +/- {post_prune_eval['ci95_reward']:.1f}') |
| next_eval_at = eval_every_n_steps |
| rollout_steps = cfg['rollout_steps'] |
| ema_interval = cfg['ema_interval'] |
| sysmon = SystemMonitor() |
| recent = deque(maxlen=100) |
| ep = 0 |
| ep_r = 0.0 |
| steps_roll = 0 |
| global_step = 0 |
| update_count = 0 |
| ep_batch_rewards = [] |
| state, _ = env.reset(seed=seed) |
| steps_to_thresh = {T: None for T in REWARD_THRESHOLDS} |
| episodes_to_thresh = {T: None for T in REWARD_THRESHOLDS} |
| solved = False |
| t_solve = None |
| ep_solve = None |
| if verbose: |
| print('=' * 80) |
| print(f'HD-PPO -> Pendulum-v1') |
| print(f' D={cfg['D']} beta={cfg['beta']} total_timesteps={total_timesteps:,}') |
| print('=' * 80) |
| t0 = time.perf_counter() |
| while global_step < total_timesteps: |
| a_raw, a_env, lp, val, H_re, H_im = agent.select_action(state) |
| next_s, reward, term, trunc, _ = env.step(a_env) |
| done = term or trunc |
| global_step += 1 |
| agent.store(H_re, H_im, a_raw, reward, lp, val, done) |
| ep_r += reward |
| steps_roll += 1 |
| if eval_csv_writer is not None and next_eval_at is not None: |
| while global_step >= next_eval_at: |
| periodic_eval = evaluate_agent(agent) |
| eval_csv_writer.writerow([next_eval_at, periodic_eval['mean_reward'], periodic_eval['ci95_reward'], 'periodic']) |
| eval_csv_file.flush() |
| if verbose: |
| print(f' [eval @ step {next_eval_at:>9,}] {periodic_eval['mean_reward']:+.1f} +/- {periodic_eval['ci95_reward']:.1f}') |
| next_eval_at += eval_every_n_steps |
| if done: |
| ep += 1 |
| recent.append(ep_r) |
| ep_batch_rewards.append(ep_r) |
| if len(recent) == 100: |
| avg100 = sum(recent) / 100.0 |
| agent.maybe_snapshot(avg100) |
| if ep % ema_interval == 0: |
| agent.ema_restore() |
| for T in REWARD_THRESHOLDS: |
| if steps_to_thresh[T] is None and avg100 >= T: |
| steps_to_thresh[T] = global_step |
| episodes_to_thresh[T] = ep |
| if not solved and avg100 >= cfg['solve_thresh']: |
| solved = True |
| t_solve = time.perf_counter() - t0 |
| ep_solve = ep |
| if verbose: |
| print(f' *** FIRST SOLVE at step {global_step:,} (ep {ep}, avg100={avg100:.1f}) -- continuing to full budget ***') |
| state, _ = env.reset() |
| ep_r = 0.0 |
| else: |
| state = next_s |
| if steps_roll >= rollout_steps: |
| t_update_start = time.perf_counter() |
| policy_loss, value_loss = agent.update(state, done) |
| update_s = time.perf_counter() - t_update_start |
| update_count += 1 |
| steps_roll = 0 |
| if csv_writer is not None and len(recent) > 0: |
| fps = rollout_steps / (update_s + 1e-08) |
| csv_writer.writerow([global_step, int(time.perf_counter() - t0), ep, len(ep_batch_rewards), float(np.mean(recent)), float(np.max(recent)), float(np.min(recent)), agent.best_avg100, policy_loss, value_loss, float(np.mean(agent.actor.log_std)), agent.entropy_coef, fps, sysmon.ram_mb()]) |
| csv_file.flush() |
| ep_batch_rewards = [] |
| if verbose and len(recent) > 0 and (update_count % 20 == 0): |
| print(f' [step {global_step:>9,}] ep {ep:>5} avg100={float(np.mean(recent)):>+8.1f} best={agent.best_avg100:>+8.1f} log_std={float(np.mean(agent.actor.log_std)):+.3f}') |
| total_time = time.perf_counter() - t0 |
| if csv_file is not None: |
| csv_file.close() |
| final_avg = sum(recent) / len(recent) if recent else 0.0 |
| eval_final = evaluate_agent(agent) |
| eval_best = None |
| if agent.best_snap is not None: |
| saved_wr, saved_wi = (agent.actor.W_re.copy(), agent.actor.W_im.copy()) |
| agent.actor.W_re, agent.actor.W_im = agent.best_snap |
| saved_log_g_ema = agent.encoder._log_g_ema |
| agent.encoder._log_g_ema = 0.0 |
| eval_best = evaluate_agent(agent) |
| agent.encoder._log_g_ema = saved_log_g_ema |
| agent.actor.W_re, agent.actor.W_im = (saved_wr, saved_wi) |
| if eval_csv_file is not None: |
| eval_csv_file.close() |
| if save_weights_path is not None: |
| use_best = eval_best is not None and eval_best['mean_reward'] > eval_final['mean_reward'] |
| W_re_save, W_im_save = agent.best_snap if use_best else (agent.actor.W_re, agent.actor.W_im) |
| np.savez(save_weights_path, W_actor_re=W_re_save, W_actor_im=W_im_save, W_critic_re=agent.critic.W_re, W_critic_im=agent.critic.W_im, critic_bias=np.float32(agent.critic.bias), fpe_phi=agent.encoder.Phi, beta_base=np.float32(cfg['beta']), feat_lo=np.array(cfg['feat_lo'], dtype=np.float32), feat_hi=np.array(cfg['feat_hi'], dtype=np.float32), D=np.int32(cfg['D']), n_feat=np.int32(agent.encoder.n_feat), action_dim=np.int32(cfg['action_dim']), log_std=agent.actor.log_std, eval_mean_final=np.float32(eval_final['mean_reward']), eval_mean_best=np.float32(eval_best['mean_reward'] if eval_best is not None else np.nan), used_best_snapshot=np.bool_(use_best)) |
| if verbose: |
| print(f' Saved actor+critic -> {save_weights_path} ({os.path.getsize(save_weights_path) / 1024:.1f} KB)') |
| if verbose: |
| eb_str = f'{eval_best['mean_reward']:+.1f}' if eval_best is not None else '-' |
| print(f'\n Training time: {total_time:.1f}s') |
| print(f' Episodes: {ep:,}') |
| print(f' Final train avg100: {final_avg:+.1f}') |
| print(f' Best train avg100: {agent.best_avg100:+.1f}') |
| print(f' Eval (final wts): {eval_final['mean_reward']:+.1f} +/- {eval_final['ci95_reward']:.1f}') |
| print(f' Eval (best wts): {eb_str}') |
| env.close() |
| return dict(seed=seed, solved=solved, total_time=total_time, final_avg100=final_avg, best_avg100=float(agent.best_avg100), eval_mean_final=eval_final['mean_reward'], eval_mean_best=eval_best['mean_reward'] if eval_best is not None else None, global_steps=global_step) |
| THIS_DIR = os.path.dirname(os.path.abspath(__file__)) |
| SEED = DEFAULT_SEED |
| STAGES = [(512, 1000000), (128, 1000000), (32, 1000000)] |
| OUT_DIR = os.path.join(THIS_DIR, 'prune_finetune') |
|
|
| def weights_path(D, stage_label): |
| return os.path.join(OUT_DIR, f'pendulum_D{D}_{stage_label}.npz') |
|
|
| def curve_csv_path(D, stage_label): |
| return os.path.join(OUT_DIR, f'training_curve_D{D}_{stage_label}.csv') |
|
|
| def eval_csv_path_for(D, stage_label): |
| return os.path.join(OUT_DIR, f'eval_curve_D{D}_{stage_label}.csv') |
|
|
| def main(): |
| os.makedirs(OUT_DIR, exist_ok=True) |
| results_json = os.path.join(OUT_DIR, 'results.json') |
| table_txt = os.path.join(OUT_DIR, 'results_table.txt') |
| warmup_jit(STAGES[0][0], PENDULUM_CONFIG['rollout_steps'], ACTION_DIM) |
| stage_records = [] |
| prev_path = None |
| t_chain0 = time.perf_counter() |
| for i, (D, timesteps) in enumerate(STAGES): |
| if i == 0: |
| stage_label = 'teacher_fresh' |
| warm_start = None |
| print(f'\n{'#' * 90}\nSTAGE {i + 1}/{len(STAGES)}: D={D} FRESH, {timesteps:,} steps\n{'#' * 90}', flush=True) |
| else: |
| stage_label = 'finetuned' |
| prev_D = STAGES[i - 1][0] |
| print(f'\n{'#' * 90}\nSTAGE {i + 1}/{len(STAGES)}: prune D={prev_D} -> D={D}, then fine-tune {timesteps:,} steps (critic warm-started)\n{'#' * 90}', flush=True) |
| prev_ckpt = np.load(prev_path) |
| warm_start = prune_actor_global(prev_ckpt, D_prime=D) |
| print(f' Pruned: kept top-{D}/{prev_D} dimensions by weight importance ({prev_D / D:.1f}x cut)') |
| path = weights_path(D, stage_label) |
| curve_csv = curve_csv_path(D, stage_label) |
| eval_csv = eval_csv_path_for(D, stage_label) |
| t0 = time.time() |
| result = train_one_seed(seed=SEED, total_timesteps=timesteps, warm_start=warm_start, save_weights_path=path, log_csv_path=curve_csv, eval_csv_path=eval_csv, eval_every_n_steps=50000, D=D if warm_start is None else None, verbose=True) |
| wall = time.time() - t0 |
| print(f' STAGE {i + 1} done: D={D} eval_final={result['eval_mean_final']:+.1f} eval_best={result['eval_mean_best']:+.1f} wall={wall:.0f}s', flush=True) |
| stage_records.append(dict(stage=stage_label, D=D, configured_timesteps=timesteps, weights_path=path, curve_csv=curve_csv, eval_csv=eval_csv, eval_mean_final=result['eval_mean_final'], eval_mean_best=result['eval_mean_best'], final_avg100=result['final_avg100'], best_avg100=result['best_avg100'], wall_time_sec=wall)) |
| with open(results_json, 'w') as f: |
| json.dump(dict(seed=SEED, stages=stage_records, complete=False), f, indent=2) |
| prev_path = path |
| total_time = time.perf_counter() - t_chain0 |
| print('\n' + '=' * 100) |
| header = f'{'stage':<16} {'D':>6} {'steps':>10} {'eval_final':>12} {'eval_best':>12} {'wall (min)':>11}' |
| print(header) |
| lines_txt = [header] |
| for r in stage_records: |
| line = f'{r['stage']:<16} {r['D']:>6} {r['configured_timesteps']:>10,} {r['eval_mean_final']:>12.1f} {r['eval_mean_best']:>12.1f} {r['wall_time_sec'] / 60:>11.1f}' |
| print(line) |
| lines_txt.append(line) |
| print(f'\nTotal chain wall time: {total_time / 60:.1f} min') |
| print('=' * 100) |
| lines_txt.append(f'\nTotal chain wall time: {total_time / 60:.1f} min') |
| with open(table_txt, 'w') as f: |
| f.write('\n'.join(lines_txt) + '\n') |
| print(f'\nWrote {table_txt}') |
| with open(results_json, 'w') as f: |
| json.dump(dict(seed=SEED, stages=stage_records, complete=True, total_chain_time_sec=total_time), f, indent=2) |
| print(f'Wrote {results_json}') |
| if __name__ == '__main__': |
| main() |
|
|
