| """Unit 8 PPO from scratch (CleanRL-style), updated for Gymnasium / Colab. |
| |
| Example: |
| python ppo.py --env-id LunarLander-v3 --repo-id while0628/ppo-LunarLander-v3-unit8 --total-timesteps 500000 --upload |
| """ |
| import argparse |
| import json |
| import os |
| import random |
| import time |
| from pathlib import Path |
|
|
| import gymnasium as gym |
| import imageio |
| import numpy as np |
| import torch |
| import torch.nn as nn |
| import torch.optim as optim |
| from torch.distributions.categorical import Categorical |
| from torch.utils.tensorboard import SummaryWriter |
|
|
|
|
| def parse_args(): |
| p = argparse.ArgumentParser(description="PPO for discrete-action Gymnasium environments") |
| p.add_argument("--exp-name", default="ppo") |
| p.add_argument("--env-id", default="LunarLander-v3") |
| p.add_argument("--repo-id", default="while0628/ppo-LunarLander-v3-unit8") |
| p.add_argument("--total-timesteps", type=int, default=500_000) |
| p.add_argument("--learning-rate", type=float, default=2.5e-4) |
| p.add_argument("--num-envs", type=int, default=4) |
| p.add_argument("--num-steps", type=int, default=1024) |
| p.add_argument("--anneal-lr", action=argparse.BooleanOptionalAction, default=True) |
| p.add_argument("--gamma", type=float, default=0.99) |
| p.add_argument("--gae-lambda", type=float, default=0.95) |
| p.add_argument("--num-minibatches", type=int, default=4) |
| p.add_argument("--update-epochs", type=int, default=4) |
| p.add_argument("--norm-adv", action=argparse.BooleanOptionalAction, default=True) |
| p.add_argument("--clip-coef", type=float, default=0.2) |
| p.add_argument("--clip-vloss", action=argparse.BooleanOptionalAction, default=True) |
| p.add_argument("--ent-coef", type=float, default=0.01) |
| p.add_argument("--vf-coef", type=float, default=0.5) |
| p.add_argument("--max-grad-norm", type=float, default=0.5) |
| p.add_argument("--target-kl", type=float, default=None) |
| p.add_argument("--seed", type=int, default=1) |
| p.add_argument("--cuda", action=argparse.BooleanOptionalAction, default=True) |
| p.add_argument("--upload", action="store_true", help="Evaluate and upload model/video to Hugging Face Hub") |
| p.add_argument("--eval-episodes", type=int, default=10) |
| args = p.parse_args() |
| args.batch_size = args.num_envs * args.num_steps |
| if args.batch_size % args.num_minibatches: |
| p.error("num-envs * num-steps must be divisible by num-minibatches") |
| args.minibatch_size = args.batch_size // args.num_minibatches |
| return args |
|
|
|
|
| def layer_init(layer, std=np.sqrt(2), bias_const=0.0): |
| nn.init.orthogonal_(layer.weight, std) |
| nn.init.constant_(layer.bias, bias_const) |
| return layer |
|
|
|
|
| class Agent(nn.Module): |
| def __init__(self, envs): |
| super().__init__() |
| obs_dim = int(np.prod(envs.single_observation_space.shape)) |
| action_dim = envs.single_action_space.n |
| self.critic = nn.Sequential( |
| layer_init(nn.Linear(obs_dim, 64)), nn.Tanh(), |
| layer_init(nn.Linear(64, 64)), nn.Tanh(), |
| layer_init(nn.Linear(64, 1), std=1.0), |
| ) |
| self.actor = nn.Sequential( |
| layer_init(nn.Linear(obs_dim, 64)), nn.Tanh(), |
| layer_init(nn.Linear(64, 64)), nn.Tanh(), |
| layer_init(nn.Linear(64, 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() |
| return action, probs.log_prob(action), probs.entropy(), self.critic(x) |
|
|
|
|
| def make_env(env_id, seed): |
| def thunk(): |
| env = gym.make(env_id) |
| env = gym.wrappers.RecordEpisodeStatistics(env) |
| env.action_space.seed(seed) |
| return env |
| return thunk |
|
|
|
|
| def evaluate_and_upload(agent, args, device, model_path, run_name): |
| from huggingface_hub import HfApi, get_token |
|
|
| if get_token() is None: |
| raise RuntimeError("No Hugging Face token found. Run `from huggingface_hub import login; login()` in Colab first.") |
| api = HfApi() |
| username = api.whoami()["name"] |
| if not args.repo_id.startswith(username + "/"): |
| raise ValueError(f"Logged in as {username}, but repo-id is {args.repo_id}") |
|
|
| output_dir = Path("hub_upload") / run_name |
| output_dir.mkdir(parents=True, exist_ok=True) |
| import shutil |
| shutil.copy2(model_path, output_dir / "model.pt") |
| shutil.copy2(__file__, output_dir / "ppo.py") |
| scores = [] |
| frames = [] |
| for episode in range(args.eval_episodes): |
| env = gym.make(args.env_id, render_mode="rgb_array") |
| obs, _ = env.reset(seed=args.seed + 10_000 + episode) |
| done = False |
| total_reward = 0.0 |
| step_count = 0 |
| while not done: |
| x = torch.as_tensor(obs, dtype=torch.float32, device=device).unsqueeze(0) |
| with torch.no_grad(): |
| action = agent.actor(x).argmax(dim=-1).item() |
| obs, reward, terminated, truncated, _ = env.step(action) |
| total_reward += float(reward) |
| done = terminated or truncated |
| if episode == 0: |
| frame = env.render() |
| if frame is not None: |
| frames.append(frame) |
| step_count += 1 |
| if step_count >= 2000: |
| break |
| env.close() |
| scores.append(total_reward) |
| print(f"Evaluation episode {episode + 1}/{args.eval_episodes}: reward={total_reward:.1f}") |
|
|
| if frames: |
| video_path = output_dir / "replay.mp4" |
| with imageio.get_writer(video_path, fps=30, codec="libx264", quality=7) as writer: |
| for frame in frames: |
| writer.append_data(frame) |
|
|
| mean_reward = float(np.mean(scores)) |
| std_reward = float(np.std(scores)) |
| (output_dir / "eval_results.json").write_text( |
| json.dumps({"env_id": args.env_id, "mean_reward": mean_reward, |
| "std_reward": std_reward, "episode_rewards": scores}, indent=2), |
| encoding="utf-8", |
| ) |
| (output_dir / "hyperparameters.json").write_text(json.dumps(vars(args), indent=2), encoding="utf-8") |
| (output_dir / "README.md").write_text( |
| f"""--- |
| license: mit |
| tags: |
| - reinforcement-learning |
| - ppo |
| - gymnasium |
| - deep-rl-course |
| model-index: |
| - name: PPO {args.env_id} |
| results: |
| - task: |
| type: reinforcement-learning |
| name: Reinforcement Learning |
| dataset: |
| name: {args.env_id} |
| type: {args.env_id} |
| metrics: |
| - type: mean_reward |
| value: {mean_reward:.3f} |
| name: Mean Reward |
| --- |
| # PPO from scratch — {args.env_id} |
| |
| Trained using a CleanRL-style PyTorch PPO implementation adapted for Gymnasium. |
| |
| Evaluation over {args.eval_episodes} episodes: **{mean_reward:.2f} ± {std_reward:.2f}**. |
| |
| - `model.pt`: actor-critic PyTorch state dict |
| - `ppo.py`: full training/evaluation/upload source |
| - `replay.mp4`: evaluation video (if rendering succeeded) |
| - `hyperparameters.json`: training settings |
| - `eval_results.json`: evaluation results |
| |
| To run locally, install PyTorch, `gymnasium[box2d]`, `imageio`, `imageio-ffmpeg`, `tensorboard`, and `huggingface_hub`. |
| """, |
| encoding="utf-8", |
| ) |
| print(f"Uploading to https://huggingface.co/{args.repo_id}") |
| api.create_repo(repo_id=args.repo_id, repo_type="model", exist_ok=True) |
| api.upload_folder(folder_path=str(output_dir), repo_id=args.repo_id, repo_type="model", |
| commit_message="Upload PPO model, evaluation and replay") |
| print(f"Upload complete: https://huggingface.co/{args.repo_id}") |
|
|
|
|
| def main(): |
| args = parse_args() |
| run_name = f"{args.env_id}__{args.exp_name}__{args.seed}__{int(time.time())}" |
| random.seed(args.seed) |
| np.random.seed(args.seed) |
| torch.manual_seed(args.seed) |
| torch.backends.cudnn.deterministic = True |
| device = torch.device("cuda" if torch.cuda.is_available() and args.cuda else "cpu") |
| writer = SummaryWriter(f"runs/{run_name}") |
| writer.add_text("hyperparameters", "```json\n" + json.dumps(vars(args), indent=2) + "\n```") |
|
|
| envs = gym.vector.SyncVectorEnv([make_env(args.env_id, args.seed + i) for i in range(args.num_envs)]) |
| if not isinstance(envs.single_action_space, gym.spaces.Discrete): |
| raise ValueError("This PPO script supports discrete action spaces only") |
| agent = Agent(envs).to(device) |
| optimizer = optim.Adam(agent.parameters(), lr=args.learning_rate, eps=1e-5) |
| obs_shape = envs.single_observation_space.shape |
| obs = torch.zeros((args.num_steps, args.num_envs) + obs_shape, device=device) |
| actions = torch.zeros((args.num_steps, args.num_envs), device=device) |
| logprobs = torch.zeros((args.num_steps, args.num_envs), device=device) |
| rewards = torch.zeros((args.num_steps, args.num_envs), device=device) |
| dones = torch.zeros((args.num_steps, args.num_envs), device=device) |
| values = torch.zeros((args.num_steps, args.num_envs), device=device) |
| global_step = 0 |
| start_time = time.time() |
| next_obs_np, _ = envs.reset(seed=args.seed) |
| next_obs = torch.as_tensor(next_obs_np, dtype=torch.float32, device=device) |
| next_done = torch.zeros(args.num_envs, device=device) |
| num_updates = max(1, args.total_timesteps // args.batch_size) |
|
|
| for update in range(1, num_updates + 1): |
| if args.anneal_lr: |
| frac = 1.0 - (update - 1.0) / num_updates |
| optimizer.param_groups[0]["lr"] = frac * args.learning_rate |
| for step in range(args.num_steps): |
| global_step += args.num_envs |
| obs[step] = next_obs |
| dones[step] = next_done |
| with torch.no_grad(): |
| action, logprob, _, value = agent.get_action_and_value(next_obs) |
| values[step] = value.flatten() |
| actions[step] = action |
| logprobs[step] = logprob |
| next_obs_np, reward, terminated, truncated, infos = envs.step(action.cpu().numpy()) |
| done = np.logical_or(terminated, truncated) |
| rewards[step] = 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.get("final_info") if infos.get("final_info") is not None else []): |
| if info and "episode" in info: |
| episode = info["episode"] |
| ep_r = float(np.asarray(episode["r"]).item()) |
| ep_l = int(np.asarray(episode["l"]).item()) |
| print(f"global_step={global_step}, episodic_return={ep_r:.2f}") |
| writer.add_scalar("charts/episodic_return", ep_r, global_step) |
| writer.add_scalar("charts/episodic_length", ep_l, global_step) |
| |
| if "episode" in infos and "_episode" in infos: |
| for i, valid in enumerate(infos["_episode"]): |
| if valid: |
| writer.add_scalar("charts/episodic_return", float(infos["episode"]["r"][i]), global_step) |
|
|
| with torch.no_grad(): |
| next_value = agent.get_value(next_obs).reshape(1, -1) |
| advantages = torch.zeros_like(rewards) |
| lastgaelam = 0 |
| for t in reversed(range(args.num_steps)): |
| if t == args.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] + args.gamma * nextvalues * nextnonterminal - values[t] |
| advantages[t] = lastgaelam = delta + args.gamma * args.gae_lambda * nextnonterminal * lastgaelam |
| returns = advantages + values |
|
|
| b_obs = obs.reshape((-1,) + obs_shape) |
| b_logprobs = logprobs.reshape(-1) |
| b_actions = actions.reshape(-1) |
| b_advantages = advantages.reshape(-1) |
| b_returns = returns.reshape(-1) |
| b_values = values.reshape(-1) |
| b_inds = np.arange(args.batch_size) |
| clipfracs = [] |
| for epoch in range(args.update_epochs): |
| np.random.shuffle(b_inds) |
| for start in range(0, args.batch_size, args.minibatch_size): |
| mb_inds = b_inds[start:start + args.minibatch_size] |
| _, newlogprob, entropy, newvalue = agent.get_action_and_value(b_obs[mb_inds], b_actions.long()[mb_inds]) |
| logratio = newlogprob - b_logprobs[mb_inds] |
| ratio = logratio.exp() |
| with torch.no_grad(): |
| old_approx_kl = (-logratio).mean() |
| approx_kl = ((ratio - 1) - logratio).mean() |
| clipfracs.append(((ratio - 1.0).abs() > args.clip_coef).float().mean().item()) |
| mb_advantages = b_advantages[mb_inds] |
| if args.norm_adv: |
| mb_advantages = (mb_advantages - mb_advantages.mean()) / (mb_advantages.std() + 1e-8) |
| pg_loss1 = -mb_advantages * ratio |
| pg_loss2 = -mb_advantages * torch.clamp(ratio, 1 - args.clip_coef, 1 + args.clip_coef) |
| pg_loss = torch.max(pg_loss1, pg_loss2).mean() |
| newvalue = newvalue.view(-1) |
| if args.clip_vloss: |
| v_loss_unclipped = (newvalue - b_returns[mb_inds]) ** 2 |
| v_clipped = b_values[mb_inds] + torch.clamp(newvalue - b_values[mb_inds], -args.clip_coef, args.clip_coef) |
| v_loss_clipped = (v_clipped - b_returns[mb_inds]) ** 2 |
| v_loss = 0.5 * torch.max(v_loss_unclipped, v_loss_clipped).mean() |
| else: |
| v_loss = 0.5 * ((newvalue - b_returns[mb_inds]) ** 2).mean() |
| entropy_loss = entropy.mean() |
| loss = pg_loss - args.ent_coef * entropy_loss + args.vf_coef * v_loss |
| optimizer.zero_grad() |
| loss.backward() |
| nn.utils.clip_grad_norm_(agent.parameters(), args.max_grad_norm) |
| optimizer.step() |
| if args.target_kl is not None and approx_kl > args.target_kl: |
| break |
| y_pred, y_true = b_values.cpu().numpy(), b_returns.cpu().numpy() |
| var_y = np.var(y_true) |
| explained_var = np.nan if var_y == 0 else 1 - np.var(y_true - y_pred) / var_y |
| writer.add_scalar("charts/learning_rate", optimizer.param_groups[0]["lr"], global_step) |
| writer.add_scalar("losses/value_loss", v_loss.item(), global_step) |
| writer.add_scalar("losses/policy_loss", pg_loss.item(), global_step) |
| writer.add_scalar("losses/entropy", entropy_loss.item(), global_step) |
| writer.add_scalar("losses/old_approx_kl", old_approx_kl.item(), global_step) |
| writer.add_scalar("losses/approx_kl", approx_kl.item(), global_step) |
| writer.add_scalar("losses/clipfrac", np.mean(clipfracs), global_step) |
| writer.add_scalar("losses/explained_variance", explained_var, global_step) |
| sps = int(global_step / max(time.time() - start_time, 1e-6)) |
| writer.add_scalar("charts/SPS", sps, global_step) |
| print(f"Update {update}/{num_updates} | Steps {global_step} | SPS {sps} | Policy loss {pg_loss.item():.4f}") |
|
|
| output_dir = Path("models") |
| output_dir.mkdir(exist_ok=True) |
| model_path = output_dir / f"{run_name}.pt" |
| torch.save(agent.state_dict(), model_path) |
| print(f"Training completed! Model saved to {model_path}") |
| envs.close() |
| writer.close() |
| if args.upload: |
| evaluate_and_upload(agent, args, device, model_path, run_name) |
| else: |
| print("To evaluate and upload, rerun with --upload (this starts a new training run).") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
