PPO.py

import argparse
import os
import time
import random
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.distributions import Categorical
from distutils.util import strtobool
from torch.utils.tensorboard import SummaryWriter
import gymnasium as gym
import wandb


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--exp-name",
        type=str,
        default=os.path.basename(__file__).rstrip(".py"),
        help="the name of this experiment",
    )
    parser.add_argument(
        "--gym-id",
        type=str,
        default="CartPole-v1",
        help="the id of the gym environment",
    )
    parser.add_argument(
        "--learning-rate",
        type=float,
        default=2.5e-4,
        help="the learning rate of the optimizer",
    )
    parser.add_argument("--seed", type=int, default=1, help="seed of the experiment")
    parser.add_argument(
        "--total-timesteps",
        type=int,
        default=200_0000,
        help="total timesteps of the experiments",
    )
    parser.add_argument(
        "--torch-deterministic",
        type=lambda x: bool(strtobool(x)),
        default=True,
        help="if toggled, `torch.backends.cudnn.deterministic=False`",
    )
    parser.add_argument(
        "--cuda",
        type=lambda x: bool(strtobool(x)),
        default=True,
        help="if toggled, cuda will be enabled by default",
    )
    parser.add_argument(
        "--track",
        type=lambda x: bool(strtobool(x)),
        default=False,
        help="if toggled, this experiment will be tracked with Weights and Biases",
    )
    parser.add_argument(
        "--wandb-project-name",
        type=str,
        default="ppo-cartpole",
        help="the wandb's project name",
    )
    parser.add_argument(
        "--wandb-entity",
        type=str,
        default=None,
        help="the entity (team) of wandb's project",
    )
    parser.add_argument(
        "--capture-video",
        type=lambda x: bool(strtobool(x)),
        default=False,
        help="whether to capture videos of the agent performances (check out `videos` folder)",
    )
    parser.add_argument(
        "--num-envs",
        type=int,
        default=4,
        help="the number of parallel game environments",
    )
    # how much data to collect per policy rollout, total num datapoints = num-steps * num-envs
    parser.add_argument(
        "--num-steps",
        type=int,
        default=128,
        help="the number of steps to run in each environment per policy rollout",
    )
    parser.add_argument(
        "--anneal-lr",
        type=lambda x: bool(strtobool(x)),
        default=True,
        help="Toggle learning rate annealing for policy and value networks",
    )
    parser.add_argument(
        "--gae",
        type=lambda x: bool(strtobool(x)),
        default=True,
        help="use generalized advantage estimation",
    )
    parser.add_argument(
        "--gae-lambda",
        type=float,
        default=0.95,
        help="the lambda for the general advantage estimation",
    )
    parser.add_argument(
        "--gamma", type=float, default=0.99, help="the discount factor gamma"
    )
    parser.add_argument(
        "--num-minibatches", type=int, default=4, help="the number of mini-batches"
    )
    parser.add_argument(
        "--update-epochs", type=int, default=4, help="the K epochs to update the policy"
    )
    parser.add_argument(
        "--norm-adv",
        type=lambda x: bool(strtobool(x)),
        default=True,
        help="Toggles advantages normalization",
    )
    parser.add_argument(
        "--clip-coeff",
        type=float,
        default=0.2,
        help="the surrogate clipping coefficient",
    )
    parser.add_argument(
        "--clip-vloss",
        type=lambda x: bool(strtobool(x)),
        default=True,
        help="Toggles whether or not to use a clipped loss for the value function, as per the paper.",
    )
    parser.add_argument(
        "--ent-coeff", type=float, default=0.01, help="coefficient of the entropy"
    )
    parser.add_argument(
        "--vf-coeff", type=float, default=0.5, help="coefficient of the value function"
    )
    parser.add_argument(
        "--max-grad-norm",
        type=float,
        default=0.5,
        help="the maximum norm for the gradient clipping",
    )
    parser.add_argument(
        "--target-kl",
        type=float,
        default=None,
        help="the target KL divergence threshold",
    )
    args = parser.parse_args()
    args.batch_size = int(args.num_envs * args.num_steps)
    args.minibatch_size = int(args.batch_size // args.num_minibatches)
    return args


def make_env(gym_id, idx, capture_video, run_name):
    def thunk():
        env = gym.make(gym_id, render_mode="rgb_array")
        env = gym.wrappers.RecordEpisodeStatistics(env)
        if capture_video and idx == 0:
            # Only capture video for the first environment
            env = gym.wrappers.RecordVideo(
                env,
                f"videos/{run_name}",
                step_trigger=lambda x: x % 100 == 0,
                video_length=0,
                episode_trigger=lambda x: True,
            )

        return env

    return thunk


def layer_init(layer, std=np.sqrt(2), bias_const=0.0):
    torch.nn.init.orthogonal_(layer.weight, std)
    torch.nn.init.constant_(layer.bias, bias_const)
    return layer


class Agent(nn.Module):
    def __init__(self, envs):
        super(Agent, self).__init__()
        self.critic = nn.Sequential(
            layer_init(
                nn.Linear(np.array(envs.single_observation_space.shape).prod(), 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(np.array(envs.single_observation_space.shape).prod(), 64)
            ),
            nn.Tanh(),
            layer_init(nn.Linear(64, 64)),
            nn.Tanh(),
            # Ensures layers params have similar scalar values
            layer_init(nn.Linear(64, envs.single_action_space.n), 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),
        )  # critc's value


if __name__ == "__main__":
    args = parse_args()
    run_name = f"{args.gym_id}__{args.exp_name}__{args.seed}__{int(time.time())}"
    writer = SummaryWriter(f"runs/{run_name}")
    writer.add_text(
        "hyperparameters",
        "|".join([f"{key}: {value}" for key, value in vars(args).items()]),
    )

    if args.track:
        import wandb

        wandb.init(
            project=args.wandb_project_name,
            entity=args.wandb_entity,
            config=vars(args),
            name=run_name,
            monitor_gym=False,
            save_code=True,
        )
        wandb.watch_called = False
        random.seed(args.seed)

    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    torch.backends.cudnn.benchmark = args.torch_deterministic
    device = torch.device("cuda" if torch.cuda.is_available() and args.cuda else "cpu")

    print(f"Experiment Name: {args.exp_name}")
    print(f"Gym ID: {args.gym_id}")
    print(f"Learning Rate: {args.learning_rate}")
    print(f"Seed: {args.seed}")
    print(f"Total Timesteps: {args.total_timesteps}")
    print(f"Torch Deterministic: {args.torch_deterministic}")
    print(f"Use CUDA: {args.cuda}")
    print("Setup complete. Ready to start training...")
    envs = gym.vector.SyncVectorEnv(
        [
            make_env(args.gym_id, i, args.capture_video, run_name)
            for i in range(args.num_envs)
        ]
    )
    assert isinstance(
        envs.single_action_space, gym.spaces.Discrete
    ), "only discrete action space is supported"
    print(
        f"Observation space: {envs.single_observation_space.shape, envs.single_observation_space}"
    )
    print(f"Action space: {envs.single_action_space.n}")
    print("Environment setup complete. Starting interaction loop...")
    agent = Agent(envs).to(device)
    print(agent)
    optimizer = optim.Adam(agent.parameters(), lr=args.learning_rate, eps=1e-5)

    # Storage for a rollout
    # Concatenate shape of observation space with (num_steps, num_envs)
    obs = torch.zeros(
        (args.num_steps, args.num_envs) + envs.single_observation_space.shape
    ).to(device)
    actions = torch.zeros(
        (args.num_steps, args.num_envs) + envs.single_action_space.shape
    ).to(device)
    logprobs = torch.zeros((args.num_steps, args.num_envs)).to(device)
    rewards = torch.zeros((args.num_steps, args.num_envs)).to(device)
    dones = torch.zeros((args.num_steps, args.num_envs)).to(device)
    values = torch.zeros((args.num_steps, args.num_envs)).to(device)

    global_step = 0
    start_time = time.time()
    # Convert the initial observation from environment reset into a PyTorch tensor and move it to specified device (CPU/GPU)
    next_obs = torch.Tensor(envs.reset()[0]).to(device)

    # Create a tensor of zeros with size equal to number of environments, representing initial "done" states
    # Each element tracks whether corresponding environment instance has finished
    next_done = torch.zeros(args.num_envs).to(device)

    # Calculate total number of updates to perform during training
    # Formula: total_timesteps / (steps_per_update * number_of_environments)
    num_updates = 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
            lrnow = frac * args.learning_rate
            optimizer.param_groups[0]["lr"] = lrnow

        # Rollout phase
        for step in range(0, args.num_steps):
            # Each step happens in parallel across the num_envs environments, hence we update global_step by num_envs
            global_step += 1 * 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, reward, terminated, truncated, info = envs.step(
                action.cpu().numpy()
            )
            rewards[step] = torch.tensor(reward).to(device).view(-1)
            next_obs = torch.Tensor(next_obs).to(device)
            next_done = torch.tensor(terminated).to(device).view(-1) | torch.tensor(
                truncated
            ).to(device).view(-1)
            if info is not None:
                # Check if 'episode' key exists in info, which holds the vectorized episode results
                if "episode" in info:
                    episode_info = info["episode"]

                    # 1. Get the vectorized data (these are likely NumPy arrays or similar)
                    episode_rewards = episode_info.get("r", None)
                    episode_lengths = episode_info.get("l", None)
                    episode_times = episode_info.get("t", None)

                    # 2. Find the index where an episode actually terminated.
                    # In vectorized environments, an episode finishes when its length is > 0.
                    # The mask identifies which parallel environment just finished an episode.
                    if episode_lengths is not None:
                        finished_mask = episode_lengths > 0

                        # 3. Iterate over all finished episodes within this batch
                        for i, finished in enumerate(finished_mask):
                            if finished:
                                # Get the scalar values for the finished episode
                                episode_reward = episode_rewards[i]
                                episode_length = episode_lengths[i]
                                episode_time_ = episode_times[i]

                                # Log the results for the completed episode
                                print(
                                    f"global_step={global_step}, episodic_return={episode_reward}, episodic_length={episode_length}, time={episode_time_}"
                                )

                                # 4. Log to TensorBoard
                                writer.add_scalar(
                                    "charts/episodic_return",
                                    episode_reward,
                                    global_step,
                                )
                                writer.add_scalar(
                                    "charts/episodic_length",
                                    episode_length,
                                    global_step,
                                )

                                # 5. Log to W&B (if tracking is enabled)
                                if args.track:
                                    wandb.log(
                                        {
                                            "charts/episodic_return": episode_reward,
                                            "charts/episodic_length": episode_length,
                                            "global_step": global_step,
                                        }
                                    )
        # Bootstrap value if using GAE
        with torch.no_grad():
            # bootstrapped value
            # The Critic network (agent.get_value) estimates the expected future return (V(st+1)) from the last state observed in the collected trajectory (next_obs). This value is used to estimate the end-of-trajectory return.
            next_value = agent.get_value(next_obs).reshape(1, -1)
            if args.gae:
                advantages = torch.zeros_like(rewards).to(device)
                lastgaelam = 0
                # reversed range iterator from num_steps-1 to 0, because we need to calculate advantages backwards
                # relies on the fact that advantages[t] depends on advantages[t+1]
                for t in reversed(range(args.num_steps)):
                    if t == args.num_steps - 1:
                        # nextnonterminal indicates if the next state is non-terminal (1.0) or terminal (0.0)
                        # next_done is a tensor indicating which environments are done after the last step
                        nextnonterminal = ~(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]
                    )
                    # 1. Calculate the current GAE term (A_t)
                    current_gaelam = (
                        delta
                        + args.gamma * args.gae_lambda * nextnonterminal * lastgaelam
                    )
                    # 2. Assign the calculated result to both the output array and the recursive variable
                    advantages[t] = current_gaelam
                    lastgaelam = current_gaelam
                returns = advantages + values
        # flatten the batch
        b_obs = obs.reshape((-1,) + envs.single_observation_space.shape)
        b_logprobs = logprobs.reshape(-1)
        b_actions = actions.reshape((-1,) + envs.single_action_space.shape)
        b_returns = returns.reshape(-1)
        b_values = values.reshape(-1)
        b_advantages = advantages.reshape(-1)

        clipfracs = []

        # Optimizing the policy and value network
        b_inds = np.arange(args.batch_size)
        for epoch in range(args.update_epochs):
            # Shuffle the indices before creating mini-batches
            np.random.shuffle(b_inds)
            # loop through the batch size(512), stepping by minibatch size (128). 4 mini-batches per epoch
            for start in range(0, args.batch_size, args.minibatch_size):
                end = start + args.minibatch_size
                # minibatch indices
                mb_inds = b_inds[start:end]
                # Evaluate the actions and values for the mini-batch
                # Add the actions so that get_action_and_value doesn't sample new actions
                _, newlogprob, entropy, newvalue = agent.get_action_and_value(
                    b_obs[mb_inds], b_actions.long()[mb_inds]
                )
                # flattensnewvalue to be a 1D tensor of shape (minibatch_size,)
                newvalue = newvalue.view(-1)
                logratio = newlogprob - b_logprobs[mb_inds]
                # get the ratio (pi_theta / pi_theta_old)
                ratio = logratio.exp()
                with torch.no_grad():
                    old_approx_kl = (-logratio).mean()
                    approx_kl = ((ratio - 1) - logratio).mean()
                    # Measures how often the policy update exceeds the clipping threshold
                    clipfracs += [
                        ((ratio - 1.0).abs() > args.clip_coeff).float().mean().item()
                    ]
                # Log the approximate KL divergence and clipping fraction

                mb_advantages = b_advantages[mb_inds]
                if args.norm_adv:
                    mb_advantages = (mb_advantages - mb_advantages.mean()) / (
                        mb_advantages.std() + 1e-8
                    )

                # Calculate the surrogate loss components
                pg_loss1 = -b_advantages[mb_inds] * ratio
                pg_loss2 = -b_advantages[mb_inds] * torch.clamp(
                    ratio, 1 - args.clip_coeff, 1 + args.clip_coeff
                )
                # max of negatives equivalent to min of positives
                pg_loss = torch.max(pg_loss1, pg_loss2).mean()

                # Value function loss with clipping
                newvalue = newvalue.view(-1)
                if args.clip_vloss:
                    v_loss_unclipped = (newvalue - b_returns[mb_inds]) ** 2
                    # similar to policy clipping, but for value function
                    v_clipped = b_values[mb_inds] + torch.clamp(
                        newvalue - b_values[mb_inds],
                        -args.clip_coeff,
                        args.clip_coeff,
                    )
                    v_loss_clipped = (v_clipped - b_returns[mb_inds]) ** 2
                    v_loss_max = torch.max(v_loss_unclipped, v_loss_clipped)
                    v_loss = 0.5 * v_loss_max.mean()
                else:
                    # usually just use MSE loss
                    v_loss = 0.5 * ((newvalue - b_returns[mb_inds]) ** 2).mean()

                # Entropy loss to encourage exploration
                entropy_loss = entropy.mean()
                # Mininize the policy loss, value loss, and maximize entropy (hence the negative sign)
                loss = pg_loss + args.vf_coeff * v_loss - args.ent_coeff * entropy_loss

                optimizer.zero_grad()
                loss.backward()
                nn.utils.clip_grad_norm_(agent.parameters(), args.max_grad_norm)
                optimizer.step()
            # Early stopping if KL divergence is too large
            if args.target_kl is not None:
                if 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
        print(f" Explained Variance: {explained_var:.2f}")
        # Logging updates
        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/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)
        print(f"SPS: {int(global_step / (time.time() - start_time))}")
        writer.add_scalar(
            "charts/SPS", int(global_step / (time.time() - start_time)), global_step
        )

    envs.close()
    # === Dedicated Evaluation and Full Video Generation ===

    # 1. Create a dedicated evaluation environment with video capture
    # Ensure the run name is distinct to avoid overwriting training videos
    eval_run_name = f"{args.gym_id}_EVAL__{args.seed}__{int(time.time())}"

    # The dedicated evaluation environment must have 'capture_video=True'
    eval_env_func = make_env(args.gym_id, 0, True, eval_run_name)
    eval_env = eval_env_func()  # Create a single, non-vectorized environment

    print("\n--- Starting Final Evaluation with Video Recording ---")

    # 2. Reset the environment
    obs, info = eval_env.reset(seed=args.seed)
    done = False
    episode_reward = 0

    # 3. Run the evaluation episode
    while not done:
        # Convert observation to tensor and move to device
        state_tensor = torch.Tensor(obs).to(device)

        # Get action from the trained agent (no_grad is essential)
        with torch.no_grad():
            action, _, _, _ = agent.get_action_and_value(state_tensor)

        # Step the environment
        obs, reward, terminated, truncated, info = eval_env.step(
            action.cpu().numpy().item()
        )
        done = terminated or truncated
        episode_reward += reward

    print(f"Evaluation Complete. Final Episode Reward: {episode_reward}")

    # 4. Close the evaluation environment to finalize the video
    eval_env.close()

    # === End of Evaluation ===
    writer.close()