train_hdppo.py

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()