Add from-scratch PPO (Unit 8): from-scratch/ppo_scratch.py

from-scratch/ppo_scratch.py

@@ -0,0 +1,154 @@
+import random, time
+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
+
+ENV_ID = "LunarLander-v2"
+NUM_ENVS = 16
+NUM_STEPS = 1024
+TOTAL_TIMESTEPS = 1000000
+LR = 2.5e-4
+GAMMA = 0.999
+GAE_LAMBDA = 0.98
+NUM_MINIBATCHES = 32
+UPDATE_EPOCHS = 4
+CLIP_COEF = 0.2
+ENT_COEF = 0.01
+VF_COEF = 0.5
+MAX_GRAD_NORM = 0.5
+SEED = 1
+
+BATCH_SIZE = NUM_ENVS * NUM_STEPS
+MINIBATCH_SIZE = BATCH_SIZE // NUM_MINIBATCHES
+NUM_UPDATES = TOTAL_TIMESTEPS // BATCH_SIZE
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print("device:", device, "updates:", NUM_UPDATES, flush=True)
+
+random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
+
+def make_env(idx):
+    def thunk():
+        env = gym.make(ENV_ID)
+        env = gym.wrappers.RecordEpisodeStatistics(env)
+        env.reset(seed=SEED + idx)
+        return env
+    return thunk
+
+envs = gym.vector.SyncVectorEnv([make_env(i) for i in range(NUM_ENVS)])
+obs_dim = int(np.array(envs.single_observation_space.shape).prod())
+act_dim = envs.single_action_space.n
+
+def layer_init(layer, std=np.sqrt(2), bias=0.0):
+    nn.init.orthogonal_(layer.weight, std)
+    nn.init.constant_(layer.bias, bias)
+    return layer
+
+class Agent(nn.Module):
+    def __init__(self):
+        super().__init__()
+        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, act_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)
+
+agent = Agent().to(device)
+optimizer = optim.Adam(agent.parameters(), lr=LR, eps=1e-5)
+
+obs = torch.zeros((NUM_STEPS, NUM_ENVS, obs_dim)).to(device)
+actions = torch.zeros((NUM_STEPS, NUM_ENVS)).to(device)
+logprobs = torch.zeros((NUM_STEPS, NUM_ENVS)).to(device)
+rewards = torch.zeros((NUM_STEPS, NUM_ENVS)).to(device)
+dones = torch.zeros((NUM_STEPS, NUM_ENVS)).to(device)
+values = torch.zeros((NUM_STEPS, NUM_ENVS)).to(device)
+
+global_step = 0
+start_time = time.time()
+next_obs, _ = envs.reset(seed=SEED)
+next_obs = torch.Tensor(next_obs).to(device)
+next_done = torch.zeros(NUM_ENVS).to(device)
+
+for update in range(1, NUM_UPDATES + 1):
+    frac = 1.0 - (update - 1.0) / NUM_UPDATES
+    optimizer.param_groups[0]["lr"] = frac * LR
+    ep_returns = []
+    for step in range(NUM_STEPS):
+        global_step += 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, term, trunc, info = envs.step(action.cpu().numpy())
+        done = np.logical_or(term, trunc)
+        rewards[step] = torch.tensor(reward).to(device).view(-1)
+        next_obs = torch.Tensor(next_obs_np).to(device)
+        next_done = torch.Tensor(done).to(device)
+        if "final_info" in info:
+            for item in info["final_info"]:
+                if item and "episode" in item:
+                    ep_returns.append(item["episode"]["r"])
+    with torch.no_grad():
+        next_value = agent.get_value(next_obs).reshape(1, -1)
+        advantages = torch.zeros_like(rewards).to(device)
+        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
+    b_obs = obs.reshape((-1, obs_dim))
+    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(BATCH_SIZE)
+    for epoch in range(UPDATE_EPOCHS):
+        np.random.shuffle(b_inds)
+        for start in range(0, BATCH_SIZE, MINIBATCH_SIZE):
+            mb_inds = b_inds[start:start + 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()
+            mb_adv = b_advantages[mb_inds]
+            mb_adv = (mb_adv - mb_adv.mean()) / (mb_adv.std() + 1e-8)
+            pg_loss1 = -mb_adv * ratio
+            pg_loss2 = -mb_adv * torch.clamp(ratio, 1 - CLIP_COEF, 1 + CLIP_COEF)
+            pg_loss = torch.max(pg_loss1, pg_loss2).mean()
+            newvalue = newvalue.view(-1)
+            v_loss = 0.5 * ((newvalue - b_returns[mb_inds]) ** 2).mean()
+            entropy_loss = entropy.mean()
+            loss = pg_loss - ENT_COEF * entropy_loss + v_loss * VF_COEF
+            optimizer.zero_grad()
+            loss.backward()
+            nn.utils.clip_grad_norm_(agent.parameters(), MAX_GRAD_NORM)
+            optimizer.step()
+    mr = np.mean(ep_returns) if ep_returns else float("nan")
+    sps = int(global_step / (time.time() - start_time))
+    print(f"update {update}/{NUM_UPDATES} step {global_step} ep_rew_mean {mr:.1f} sps {sps}", flush=True)
+
+envs.close()
+torch.save(agent.state_dict(), "ppo_scratch_lunarlander.pt")
+print("=== TRAINING DONE, weights saved ===", flush=True)