| |
| """ |
| PPO 训练脚本 —— LunarLander（第 4 周 · 第一个智能体） |
| ===================================================== |
| 对应课件《第4周-强化学习进阶实践内容》第 2~14 页：用 PPO 训练一个月球着陆器智能体。 |
| |
| 与第一周（Stable-Baselines3 现成库）不同，本脚本**从零手写 PPO 的全部核心组件**， |
| 每个组件都能在代码里直接看到： |
| |
| 1. Actor-Critic 网络 —— get_action_and_value() |
| 2. rollout 采样 —— 与环境交互收集经验 |
| 3. GAE 广义优势估计 —— compute_gae() |
| 4. PPO-Clip 损失 —— 策略/价值/熵三项加权 |
| 5. 评估、存盘、视频录制、日志 |
| |
| 快速开始 |
| -------- |
| 冒烟测试（约 20 秒）:: |
| |
| python ppo.py --total-timesteps 16384 --num-envs 8 --eval-interval 5 \ |
| --eval-episodes 4 --out-dir runs --run-name smoke |
| |
| 正式训练（CPU 约 15~20 分钟，可达到官方 "solved" 标准 200 分）:: |
| |
| python ppo.py --total-timesteps 2500000 --num-envs 16 --num-steps 256 \ |
| --gamma 0.999 --gae-lambda 0.98 --num-minibatches 8 --update-epochs 10 \ |
| --eval-interval 15 --eval-episodes 16 --capture-video \ |
| --out-dir runs --run-name main_2.5m |
| |
| 复现课件命令（5 万步，会失败——这正是课件想演示的"步数不够"）:: |
| |
| python ppo.py --total-timesteps 50000 --num-envs 4 --num-steps 128 \ |
| --out-dir runs --run-name baseline_50k |
| """ |
|
|
| import argparse |
| import csv |
| import json |
| import os |
| import random |
| import time |
| from collections import deque |
|
|
| import numpy as np |
| import torch |
| import torch.nn as nn |
| import torch.optim as optim |
| from torch.distributions.categorical import Categorical |
|
|
| import gymnasium as gym |
|
|
|
|
| |
| |
| |
| def layer_init(layer, std=np.sqrt(2), bias_const=0.0): |
| """正交初始化（PPO 的常规做法，训练更稳定）。""" |
| nn.init.orthogonal_(layer.weight, std) |
| nn.init.constant_(layer.bias, bias_const) |
| return layer |
|
|
|
|
| def make_env(env_id, seed, idx, render_mode=None): |
| """创建一个环境；用 seed+idx 保证每个并行环境轨迹不同但可复现。""" |
| return gym.make(env_id, render_mode=render_mode) |
|
|
|
|
| |
| |
| |
| class VecEnv: |
| """N 个环境并排前进（gymnasium API：reset -> (obs, info), step -> 5 元组）。""" |
|
|
| def __init__(self, envs): |
| self.envs = envs |
| self.num_envs = len(envs) |
| self.episode_returns = deque(maxlen=200) |
| self.episode_lengths = deque(maxlen=200) |
| self._cur_ret = np.zeros(self.num_envs, dtype=np.float64) |
| self._cur_len = np.zeros(self.num_envs, dtype=np.int64) |
| self.obs_dim = int(np.prod(envs[0].observation_space.shape)) |
| self.action_dim = int(envs[0].action_space.n) |
|
|
| def reset(self, seeds=None): |
| obs = np.zeros((self.num_envs, self.obs_dim), dtype=np.float32) |
| for i, env in enumerate(self.envs): |
| seed = None if seeds is None else int(seeds[i]) |
| o, _ = env.reset(seed=seed) |
| obs[i] = np.asarray(o, dtype=np.float32).ravel() |
| self._cur_ret[i] = 0.0 |
| self._cur_len[i] = 0 |
| return obs |
|
|
| def step(self, actions): |
| """同时推进一步；回合结束后自动 reset（不重新播种，保留随机性）。""" |
| obs = np.zeros((self.num_envs, self.obs_dim), dtype=np.float32) |
| rewards = np.zeros(self.num_envs, dtype=np.float32) |
| dones = np.zeros(self.num_envs, dtype=np.float32) |
| infos = [] |
| for i, env in enumerate(self.envs): |
| o, r, terminated, truncated, info = env.step(int(actions[i])) |
| done = bool(terminated or truncated) |
| self._cur_ret[i] += r |
| self._cur_len[i] += 1 |
| info = dict(info) if isinstance(info, dict) else {} |
| if done: |
| info["episode_return"] = float(self._cur_ret[i]) |
| info["episode_length"] = int(self._cur_len[i]) |
| self.episode_returns.append(info["episode_return"]) |
| self.episode_lengths.append(info["episode_length"]) |
| self._cur_ret[i] = 0.0 |
| self._cur_len[i] = 0 |
| o, _ = env.reset() |
| obs[i] = np.asarray(o, dtype=np.float32).ravel() |
| rewards[i] = r |
| dones[i] = float(done) |
| infos.append(info) |
| return obs, rewards, dones, infos |
|
|
| def close(self): |
| for env in self.envs: |
| try: |
| env.close() |
| except Exception: |
| pass |
|
|
|
|
| |
| |
| |
| class VideoRecorder: |
| def __init__(self, path, fps=30): |
| self.path = path |
| self.fps = fps |
| self.frames = [] |
|
|
| def add(self, frame): |
| if frame is not None: |
| self.frames.append(np.asarray(frame, dtype=np.uint8)) |
|
|
| def save(self): |
| if not self.frames: |
| return None |
| import imageio |
|
|
| os.makedirs(os.path.dirname(os.path.abspath(self.path)), exist_ok=True) |
| writer = imageio.get_writer(self.path, fps=self.fps, codec="libx264", |
| quality=8, macro_block_size=1) |
| for frame in self.frames: |
| writer.append_data(frame) |
| writer.close() |
| return self.path |
|
|
|
|
| |
| |
| |
| def build_agent(obs_dim, action_dim, hidden=64): |
| class Agent(nn.Module): |
| def __init__(self): |
| super().__init__() |
| |
| self.critic = nn.Sequential( |
| layer_init(nn.Linear(obs_dim, hidden)), nn.Tanh(), |
| layer_init(nn.Linear(hidden, hidden)), nn.Tanh(), |
| layer_init(nn.Linear(hidden, 1), std=1.0), |
| ) |
| |
| self.actor = nn.Sequential( |
| layer_init(nn.Linear(obs_dim, hidden)), nn.Tanh(), |
| layer_init(nn.Linear(hidden, hidden)), nn.Tanh(), |
| layer_init(nn.Linear(hidden, action_dim), std=0.01), |
| ) |
|
|
| def get_value(self, x): |
| return self.critic(x) |
|
|
| def get_action_and_value(self, x, action=None): |
| """给定状态 -> (动作, 对数概率, 熵, 状态价值)。""" |
| logits = self.actor(x) |
| probs = Categorical(logits=logits) |
| if action is None: |
| action = probs.sample() |
| value = self.critic(x).view(-1) |
| return action, probs.log_prob(action), probs.entropy(), value |
|
|
| return Agent() |
|
|
|
|
| |
| |
| |
| def compute_gae(rewards, values, dones, next_done, next_value, gamma, gae_lambda, num_steps): |
| advantages = torch.zeros_like(rewards) |
| lastgaelam = 0 |
| for t in reversed(range(num_steps)): |
| if t == num_steps - 1: |
| nextnonterminal = 1.0 - next_done |
| nextvalues = next_value |
| else: |
| nextnonterminal = 1.0 - dones[t + 1] |
| nextvalues = values[t + 1] |
| |
| delta = rewards[t] + gamma * nextvalues * nextnonterminal - values[t] |
| |
| advantages[t] = lastgaelam = delta + gamma * gae_lambda * nextnonterminal * lastgaelam |
| returns = advantages + values |
| return advantages, returns |
|
|
|
|
| |
| |
| |
| @torch.no_grad() |
| def evaluate(agent, env_id, device, episodes=10, deterministic=False, seed=12345): |
| """跑若干完整回合，返回 (回报列表, 回合长度列表)。""" |
| env = gym.make(env_id) |
| returns, lengths = [], [] |
| for ep in range(episodes): |
| obs, _ = env.reset(seed=seed + ep) |
| done, ret, ln = False, 0.0, 0 |
| while not done: |
| obs_t = torch.as_tensor(obs, dtype=torch.float32, device=device).unsqueeze(0) |
| if deterministic: |
| action = int(torch.argmax(agent.actor(obs_t), dim=-1).item()) |
| else: |
| action = int(agent.get_action_and_value(obs_t)[0].item()) |
| obs, r, terminated, truncated, _ = env.step(action) |
| ret += float(r) |
| ln += 1 |
| done = bool(terminated or truncated) |
| returns.append(ret) |
| lengths.append(ln) |
| env.close() |
| return returns, lengths |
|
|
|
|
| @torch.no_grad() |
| def record_episode(agent, env_id, device, path, seed=999, deterministic=True, fps=30): |
| """录制一个完整回合为 mp4，返回 {episode_return, episode_length, frames, video}。""" |
| env = gym.make(env_id, render_mode="rgb_array") |
| obs, _ = env.reset(seed=seed) |
| rec = VideoRecorder(path, fps=fps) |
| done, ret, ln = False, 0.0, 0 |
| while not done and ln < 4000: |
| frame = env.render() |
| if frame is not None: |
| rec.add(frame) |
| obs_t = torch.as_tensor(obs, dtype=torch.float32, device=device).unsqueeze(0) |
| if deterministic: |
| action = int(torch.argmax(agent.actor(obs_t), dim=-1).item()) |
| else: |
| action = int(agent.get_action_and_value(obs_t)[0].item()) |
| obs, r, terminated, truncated, _ = env.step(action) |
| ret += float(r) |
| ln += 1 |
| done = bool(terminated or truncated) |
| env.close() |
| rec.save() |
| return {"episode_return": ret, "episode_length": ln, "frames": len(rec.frames), "video": path} |
|
|
|
|
| def summarize(returns): |
| arr = np.asarray(returns, dtype=np.float64) |
| return { |
| "episodes": int(arr.size), |
| "mean": float(arr.mean()), |
| "std": float(arr.std()), |
| "min": float(arr.min()), |
| "max": float(arr.max()), |
| "median": float(np.median(arr)), |
| "landing_rate": float(np.mean(arr >= 200.0)), |
| "p25": float(np.percentile(arr, 25)), |
| "p75": float(np.percentile(arr, 75)), |
| } |
|
|
|
|
| |
| |
| |
| def parse_args(): |
| p = argparse.ArgumentParser(description="从零手写 PPO 训练 LunarLander") |
| |
| p.add_argument("--env-id", type=str, default="LunarLander-v3", |
| help="环境 ID（gymnasium 中 LunarLander 只有 v3）") |
| p.add_argument("--seed", type=int, default=1, help="随机种子（默认 1）") |
| p.add_argument("--total-timesteps", type=int, default=1_000_000, help="总训练步数") |
| p.add_argument("--out-dir", type=str, default="runs", help="结果根目录") |
| p.add_argument("--run-name", type=str, default="", help="本次实验名（默认自动生成）") |
| p.add_argument("--no-upload", action="store_true", default=False, |
| help="不上传 HuggingFace（上传用 upload_to_hf.py）") |
| p.add_argument("--repo-id", type=str, default=None, help="HuggingFace 仓库 ID（可选）") |
| |
| p.add_argument("--num-envs", type=int, default=8, help="并行环境数（默认 8）") |
| p.add_argument("--num-steps", type=int, default=128, help="每个环境每次采样多少步") |
| p.add_argument("--capture-video", action="store_true", help="录制训练过程中的学习视频") |
| |
| p.add_argument("--learning-rate", type=float, default=2.5e-4, help="学习率（默认 2.5e-4）") |
| p.add_argument("--anneal-lr", action="store_true", default=True, help="学习率线性衰减") |
| p.add_argument("--no-anneal-lr", dest="anneal_lr", action="store_false") |
| p.add_argument("--gamma", type=float, default=0.99, help="折扣因子") |
| p.add_argument("--gae-lambda", type=float, default=0.95, help="GAE lambda") |
| p.add_argument("--num-minibatches", type=int, default=4, help="mini-batch 数量") |
| p.add_argument("--update-epochs", type=int, default=4, help="每个 rollout 更新几次") |
| p.add_argument("--norm-adv", action="store_true", default=True, help="优势标准化") |
| p.add_argument("--no-norm-adv", dest="norm_adv", action="store_false") |
| p.add_argument("--clip-coef", type=float, default=0.2, help="PPO 裁剪系数") |
| p.add_argument("--clip-vloss", action="store_true", default=True, help="价值函数裁剪") |
| p.add_argument("--no-clip-vloss", dest="clip_vloss", action="store_false") |
| p.add_argument("--ent-coef", type=float, default=0.01, help="熵奖励系数") |
| p.add_argument("--vf-coef", type=float, default=0.5, help="价值损失系数") |
| p.add_argument("--max-grad-norm", type=float, default=0.5, help="梯度裁剪范数") |
| p.add_argument("--hidden", type=int, default=64, help="隐藏层宽度") |
| |
| p.add_argument("--eval-interval", type=int, default=20, help="每多少次 update 评估一次") |
| p.add_argument("--eval-episodes", type=int, default=10, help="每次评估跑多少回合") |
| p.add_argument("--eval-greedy", action="store_true", default=True, |
| help="评估用 argmax（部署策略）；训练采样仍用随机策略") |
| p.add_argument("--eval-stochastic", dest="eval_greedy", action="store_false") |
| p.add_argument("--save-interval", type=int, default=50, help="每多少次 update 存一次盘") |
| p.add_argument("--log-interval", type=int, default=10, help="控制台打印间隔") |
| p.add_argument("--torch-threads", type=int, default=8, help="torch 线程数") |
| return p.parse_args() |
|
|
|
|
| |
| |
| |
| def main(): |
| args = parse_args() |
| torch.set_num_threads(args.torch_threads) |
| random.seed(args.seed) |
| np.random.seed(args.seed) |
| torch.manual_seed(args.seed) |
|
|
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| run_name = args.run_name or f"PPO-{args.env_id}-{int(time.time())}" |
| run_dir = os.path.join(args.out_dir, run_name) |
| os.makedirs(run_dir, exist_ok=True) |
| with open(os.path.join(run_dir, "config.json"), "w", encoding="utf-8") as f: |
| json.dump(vars(args), f, indent=2, ensure_ascii=False) |
|
|
| |
| envs = [make_env(args.env_id, args.seed, i) for i in range(args.num_envs)] |
| vec = VecEnv(envs) |
| obs_dim, action_dim = vec.obs_dim, vec.action_dim |
| agent = build_agent(obs_dim, action_dim, args.hidden).to(device) |
| optimizer = optim.Adam(agent.parameters(), lr=args.learning_rate, eps=1e-5) |
|
|
| batch_size = args.num_envs * args.num_steps |
| minibatch_size = batch_size // args.num_minibatches |
| num_iterations = max(1, args.total_timesteps // batch_size) |
|
|
| print(f"[配置] 环境={args.env_id} 设备={device} 并行环境={args.num_envs} " |
| f"每轮步数={batch_size} 迭代={num_iterations} 总步数={args.total_timesteps}") |
| print(f"[配置] obs_dim={obs_dim} action_dim={action_dim} minibatch={minibatch_size}") |
|
|
| |
| obs_buf = torch.zeros((args.num_steps, args.num_envs, obs_dim), dtype=torch.float32, device=device) |
| act_buf = torch.zeros((args.num_steps, args.num_envs), dtype=torch.float32, device=device) |
| logp_buf = torch.zeros((args.num_steps, args.num_envs), dtype=torch.float32, device=device) |
| rew_buf = torch.zeros((args.num_steps, args.num_envs), dtype=torch.float32, device=device) |
| done_buf = torch.zeros((args.num_steps, args.num_envs), dtype=torch.float32, device=device) |
| val_buf = torch.zeros((args.num_steps, args.num_envs), dtype=torch.float32, device=device) |
|
|
| next_obs = torch.as_tensor(vec.reset(seeds=[args.seed + i for i in range(args.num_envs)]), |
| dtype=torch.float32, device=device) |
| next_done = torch.zeros(args.num_envs, dtype=torch.float32, device=device) |
|
|
| |
| train_csv = open(os.path.join(run_dir, "train_log.csv"), "w", newline="", encoding="utf-8") |
| train_writer = csv.writer(train_csv) |
| train_writer.writerow(["iteration", "global_step", "sps", "episodes_in_window", |
| "mean_return", "std_return", "min_return", "max_return", |
| "mean_ep_length", "policy_loss", "value_loss", "entropy", |
| "approx_kl", "clipfrac", "learning_rate"]) |
| eval_csv = open(os.path.join(run_dir, "eval_log.csv"), "w", newline="", encoding="utf-8") |
| eval_writer = csv.writer(eval_csv) |
| eval_writer.writerow(["iteration", "global_step", "eval_mean", "eval_std", |
| "eval_min", "eval_max", "landing_rate"]) |
|
|
| video_points = set() |
| if args.capture_video: |
| video_points = {max(1, int(num_iterations * p)) for p in (0.1, 0.25, 0.4, 0.55, 0.7, 0.85, 1.0)} |
|
|
| global_step = 0 |
| t0 = time.time() |
| best_eval = -float("inf") |
| recent_returns = deque(maxlen=100) |
|
|
| for iteration in range(1, num_iterations + 1): |
| |
| if args.anneal_lr: |
| frac = 1.0 - (iteration - 1.0) / num_iterations |
| lr_now = frac * args.learning_rate |
| optimizer.param_groups[0]["lr"] = lr_now |
| else: |
| lr_now = args.learning_rate |
|
|
| |
| step_returns = [] |
| for t in range(args.num_steps): |
| global_step += args.num_envs |
| obs_buf[t] = next_obs |
| done_buf[t] = next_done |
| with torch.no_grad(): |
| action, logprob, _, value = agent.get_action_and_value(next_obs) |
| val_buf[t] = value.flatten() |
| act_buf[t] = action.float() |
| logp_buf[t] = logprob |
| next_obs_np, reward, done, infos = vec.step(action.cpu().numpy()) |
| rew_buf[t] = torch.as_tensor(reward, dtype=torch.float32, device=device) |
| next_obs = torch.as_tensor(next_obs_np, dtype=torch.float32, device=device) |
| next_done = torch.as_tensor(done, dtype=torch.float32, device=device) |
| for info in infos: |
| if "episode_return" in info: |
| step_returns.append(info["episode_return"]) |
| recent_returns.append(info["episode_return"]) |
|
|
| |
| with torch.no_grad(): |
| next_value = agent.get_value(next_obs).reshape(1, -1) |
| advantages, returns = compute_gae( |
| rew_buf, val_buf, done_buf, next_done, next_value, |
| args.gamma, args.gae_lambda, args.num_steps, |
| ) |
| b_obs = obs_buf.reshape((-1, obs_dim)) |
| b_logprobs = logp_buf.reshape(-1) |
| b_actions = act_buf.reshape(-1) |
| b_advantages = advantages.reshape(-1) |
| b_returns = returns.reshape(-1) |
| b_values = val_buf.reshape(-1) |
|
|
| |
| b_inds = np.arange(batch_size) |
| clipfracs = [] |
| pg_loss = torch.tensor(0.0) |
| v_loss = torch.tensor(0.0) |
| entropy = torch.tensor(0.0) |
| approx_kl = torch.tensor(0.0) |
| for _ in range(args.update_epochs): |
| np.random.shuffle(b_inds) |
| for start in range(0, batch_size, minibatch_size): |
| mb = b_inds[start:start + minibatch_size] |
| _, newlogprob, newentropy, newvalue = agent.get_action_and_value( |
| b_obs[mb], b_actions.long()[mb]) |
| logratio = newlogprob - b_logprobs[mb] |
| ratio = logratio.exp() |
| with torch.no_grad(): |
| approx_kl = ((ratio - 1) - logratio).mean() |
| clipfracs.append(((ratio - 1.0).abs() > args.clip_coef).float().mean().item()) |
|
|
| mb_adv = b_advantages[mb] |
| if args.norm_adv: |
| mb_adv = (mb_adv - mb_adv.mean()) / (mb_adv.std() + 1e-8) |
|
|
| |
| pg1 = -mb_adv * ratio |
| pg2 = -mb_adv * torch.clamp(ratio, 1 - args.clip_coef, 1 + args.clip_coef) |
| pg_loss = torch.max(pg1, pg2).mean() |
|
|
| |
| newvalue = newvalue.view(-1) |
| if args.clip_vloss: |
| v_unclipped = (newvalue - b_returns[mb]) ** 2 |
| v_clipped = b_values[mb] + torch.clamp( |
| newvalue - b_values[mb], -args.clip_coef, args.clip_coef) |
| v_loss = 0.5 * torch.max(v_unclipped, (v_clipped - b_returns[mb]) ** 2).mean() |
| else: |
| v_loss = 0.5 * ((newvalue - b_returns[mb]) ** 2).mean() |
|
|
| |
| entropy = newentropy.mean() |
|
|
| |
| loss = pg_loss - args.ent_coef * entropy + args.vf_coef * v_loss |
| optimizer.zero_grad() |
| loss.backward() |
| nn.utils.clip_grad_norm_(agent.parameters(), args.max_grad_norm) |
| optimizer.step() |
|
|
| |
| sps = int(global_step / max(1e-9, time.time() - t0)) |
| mean_ret = float(np.mean(step_returns)) if step_returns else float("nan") |
| std_ret = float(np.std(step_returns)) if step_returns else float("nan") |
| min_ret = float(np.min(step_returns)) if step_returns else float("nan") |
| max_ret = float(np.max(step_returns)) if step_returns else float("nan") |
| mean_len = float(np.mean(vec.episode_lengths)) if len(vec.episode_lengths) else float("nan") |
| train_writer.writerow([ |
| iteration, global_step, sps, len(step_returns), |
| f"{mean_ret:.3f}", f"{std_ret:.3f}", f"{min_ret:.3f}", f"{max_ret:.3f}", |
| f"{mean_len:.1f}", f"{float(pg_loss.detach()):.5f}", f"{float(v_loss.detach()):.5f}", |
| f"{float(entropy.detach()):.5f}", f"{float(approx_kl.detach()):.5f}", |
| f"{float(np.mean(clipfracs)):.4f}", f"{lr_now:.6g}", |
| ]) |
| train_csv.flush() |
|
|
| |
| eval_stats = None |
| if iteration % args.eval_interval == 0 or iteration == num_iterations: |
| rets, _ = evaluate(agent, args.env_id, device, |
| episodes=args.eval_episodes, |
| deterministic=args.eval_greedy, |
| seed=int(args.seed * 1000 + iteration)) |
| eval_stats = summarize(rets) |
| eval_writer.writerow([ |
| iteration, global_step, |
| f"{eval_stats['mean']:.2f}", f"{eval_stats['std']:.2f}", |
| f"{eval_stats['min']:.2f}", f"{eval_stats['max']:.2f}", |
| f"{eval_stats['landing_rate']:.3f}", |
| ]) |
| eval_csv.flush() |
| if eval_stats["mean"] > best_eval: |
| best_eval = eval_stats["mean"] |
| torch.save({ |
| "state_dict": agent.state_dict(), "obs_dim": obs_dim, |
| "action_dim": action_dim, "hidden": args.hidden, |
| "iteration": iteration, "global_step": global_step, |
| "eval_mean": eval_stats["mean"], "config": vars(args), |
| "eval_stats": eval_stats, |
| }, os.path.join(run_dir, "model_best.pth")) |
|
|
| if iteration % args.log_interval == 0 or iteration == 1 or iteration == num_iterations: |
| ev = "" if eval_stats is None else f" | eval={eval_stats['mean']:.1f}" |
| print(f"[{iteration:4d}/{num_iterations}] step={global_step:8d} " |
| f"train_ret={mean_ret:8.2f} ep={len(step_returns):3d} " |
| f"len={mean_len:6.1f} pg={float(pg_loss.detach()):7.4f} " |
| f"vl={float(v_loss.detach()):7.4f} " |
| f"ent={float(entropy.detach()):6.3f} sps={sps:5d}{ev}", flush=True) |
|
|
| |
| if iteration % args.save_interval == 0 or iteration == num_iterations: |
| torch.save({ |
| "state_dict": agent.state_dict(), "obs_dim": obs_dim, |
| "action_dim": action_dim, "hidden": args.hidden, |
| "iteration": iteration, "global_step": global_step, |
| "config": vars(args), |
| }, os.path.join(run_dir, "model_ckpt.pth")) |
|
|
| |
| if args.capture_video and iteration in video_points: |
| path = os.path.join(run_dir, "videos", f"learning_{global_step:08d}.mp4") |
| try: |
| info = record_episode(agent, args.env_id, device, path, |
| seed=args.seed * 7 + iteration, deterministic=True) |
| print(f" [视频] step={global_step} 回报={info['episode_return']:.1f} -> {path}", flush=True) |
| except Exception as e: |
| print(f" [视频] 失败: {e}", flush=True) |
|
|
| |
| torch.save({ |
| "state_dict": agent.state_dict(), "obs_dim": obs_dim, |
| "action_dim": action_dim, "hidden": args.hidden, |
| "iteration": num_iterations, "global_step": global_step, |
| "config": vars(args), |
| }, os.path.join(run_dir, "model_final.pth")) |
|
|
| final_rets, _ = evaluate(agent, args.env_id, device, |
| episodes=min(50, max(10, args.eval_episodes)), |
| deterministic=True, seed=args.seed * 999 + 7) |
| summary = { |
| "run_name": run_name, |
| "run_dir": run_dir, |
| "env_id": args.env_id, |
| "seed": args.seed, |
| "total_timesteps": args.total_timesteps, |
| "num_envs": args.num_envs, |
| "num_steps": args.num_steps, |
| "gamma": args.gamma, |
| "gae_lambda": args.gae_lambda, |
| "num_minibatches": args.num_minibatches, |
| "update_epochs": args.update_epochs, |
| "learning_rate": args.learning_rate, |
| "device": str(device), |
| "train_seconds": round(time.time() - t0, 1), |
| "final_eval": summarize(final_rets), |
| "best_eval_mean": best_eval, |
| "train_mean_return_last100": float(np.mean(recent_returns)) if recent_returns else None, |
| } |
| with open(os.path.join(run_dir, "summary.json"), "w", encoding="utf-8") as f: |
| json.dump(summary, f, indent=2, ensure_ascii=False) |
| train_csv.close() |
| eval_csv.close() |
| vec.close() |
|
|
| print(f"\n[完成] 用时 {summary['train_seconds']/60:.1f} 分钟 | " |
| f"最终评估({summary['final_eval']['episodes']}局,argmax) = " |
| f"{summary['final_eval']['mean']:.2f} ± {summary['final_eval']['std']:.2f} | " |
| f"着陆率 {summary['final_eval']['landing_rate']*100:.0f}%") |
| print(f"[完成] 结果目录: {run_dir}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
