| import argparse |
| import os |
| import gymnasium as gym |
| import numpy as np |
| import train_hdppo as m |
|
|
| def _warmup_encoder(agent, env_id, n_steps=200): |
| env = gym.make(env_id) |
| obs, _ = env.reset(seed=0) |
| for i in range(n_steps): |
| Hr, Hi = agent.encoder.encode(obs) |
| action = agent.actor.greedy(Hr, Hi) |
| obs, _, term, trunc, _ = env.step(action) |
| if term or trunc: |
| obs, _ = env.reset(seed=i + 1) |
| env.close() |
|
|
| def load_agent_from_checkpoint(path, seed=42, warmup=True): |
| data = np.load(path) |
| cfg = dict(m.PENDULUM_CONFIG) |
| cfg['D'] = int(data['D']) |
| cfg['beta'] = float(data['beta_base']) |
| cfg['fpe_phi_init'] = data['fpe_phi'] |
| if 'feat_lo' in data: |
| cfg['feat_lo'] = data['feat_lo'].tolist() |
| cfg['feat_hi'] = data['feat_hi'].tolist() |
| agent = m.HDPPOAgentContinuous(cfg, seed=seed) |
| agent.actor.W_re = data['W_actor_re'].copy() |
| agent.actor.W_im = data['W_actor_im'].copy() |
| if 'log_std' in data: |
| agent.actor.log_std = data['log_std'].copy() |
| if 'W_critic_re' in data: |
| agent.critic.W_re = data['W_critic_re'].copy() |
| agent.critic.W_im = data['W_critic_im'].copy() |
| agent.critic.bias = float(data['critic_bias']) |
| if warmup: |
| _warmup_encoder(agent, 'Pendulum-v1') |
| return agent |
|
|
| def resolve_weights(weights_arg): |
| if os.path.isfile(weights_arg): |
| return weights_arg |
| try: |
| from huggingface_hub import hf_hub_download |
| except ImportError as exc: |
| raise SystemExit('Install huggingface_hub to load remote checkpoints: pip install huggingface_hub') from exc |
| return hf_hub_download(repo_id=weights_arg, filename='hdppo-Pendulum-v1/weights.npz') |
|
|
| def main(): |
| parser = argparse.ArgumentParser(description='Enjoy HD-PPO on Pendulum-v1') |
| parser.add_argument('--weights', default='hdppo-Pendulum-v1/weights.npz') |
| parser.add_argument('--episodes', type=int, default=5) |
| parser.add_argument('--seed', type=int, default=10000) |
| parser.add_argument('--render', action='store_true') |
| args = parser.parse_args() |
| weights = resolve_weights(args.weights) |
| agent = load_agent_from_checkpoint(weights, seed=args.seed) |
| if args.render: |
| env = gym.make('Pendulum-v1', render_mode='human') |
| for ep in range(args.episodes): |
| obs, _ = env.reset(seed=args.seed + ep) |
| done = False |
| ep_r = 0.0 |
| while not done: |
| Hr, Hi = agent.encoder.encode(obs) |
| action = agent.actor.greedy(Hr, Hi) |
| obs, reward, term, trunc, _ = env.step(action) |
| ep_r += reward |
| done = term or trunc |
| print(f'episode {ep + 1}: return={ep_r:.1f}') |
| env.close() |
| else: |
| stats = m.evaluate_agent(agent, n_episodes=args.episodes, seed_base=args.seed) |
| print(f'mean={stats['mean_reward']:.2f} ci95={stats['ci95_reward']:.2f}') |
| if __name__ == '__main__': |
| main() |
|
|
