| import argparse
|
| import os
|
| import random
|
| import time
|
| from distutils.util import strtobool
|
|
|
| import gym
|
| 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
|
| from huggingface_hub import HfApi, create_repo
|
|
|
| def parse_args():
|
| parser = argparse.ArgumentParser()
|
| parser.add_argument("--exp-name", type=str, default=os.path.basename(__file__).rstrip(".py"))
|
| parser.add_argument("--env-id", type=str, default="LunarLander-v2")
|
| parser.add_argument("--seed", type=int, default=1)
|
| parser.add_argument("--torch-deterministic", type=lambda x: bool(strtobool(x)), default=True)
|
| parser.add_argument("--cuda", type=lambda x: bool(strtobool(x)), default=True)
|
| parser.add_argument("--capture-video", type=lambda x: bool(strtobool(x)), default=True)
|
|
|
|
|
| parser.add_argument("--total-timesteps", type=int, default=2000000)
|
| parser.add_argument("--learning-rate", type=float, default=3e-4)
|
| parser.add_argument("--num-envs", type=int, default=16)
|
| parser.add_argument("--num-steps", type=int, default=1024)
|
| parser.add_argument("--gamma", type=float, default=0.999)
|
| parser.add_argument("--gae-lambda", type=float, default=0.98)
|
| parser.add_argument("--num-minibatches", type=int, default=256)
|
| parser.add_argument("--update-epochs", type=int, default=4)
|
| parser.add_argument("--clip-coef", type=float, default=0.2)
|
| parser.add_argument("--ent-coef", type=float, default=0.01)
|
| parser.add_argument("--vf-coef", type=float, default=0.5)
|
| parser.add_argument("--max-grad-norm", type=float, default=0.5)
|
|
|
| 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(env_id, seed, idx, capture_video, run_name):
|
| def thunk():
|
| env = gym.make(env_id)
|
| env = gym.wrappers.RecordEpisodeStatistics(env)
|
| if capture_video and idx == 0:
|
| env = gym.wrappers.RecordVideo(env, f"videos/{run_name}")
|
| env.seed(seed)
|
| env.action_space.seed(seed)
|
| env.observation_space.seed(seed)
|
| 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().__init__()
|
| input_size = np.array(envs.single_observation_space.shape).prod()
|
| self.critic = nn.Sequential(
|
| layer_init(nn.Linear(input_size, 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(input_size, 64)),
|
| nn.Tanh(),
|
| layer_init(nn.Linear(64, 64)),
|
| nn.Tanh(),
|
| 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)
|
|
|
| if __name__ == "__main__":
|
| args = parse_args()
|
| run_name = f"{args.env_id}__{args.exp_name}__{args.seed}__{int(time.time())}"
|
| writer = SummaryWriter(f"runs/{run_name}")
|
|
|
| random.seed(args.seed)
|
| np.random.seed(args.seed)
|
| torch.manual_seed(args.seed)
|
| device = torch.device("cuda" if torch.cuda.is_available() and args.cuda else "cpu")
|
|
|
| envs = gym.vector.SyncVectorEnv(
|
| [make_env(args.env_id, args.seed + i, i, args.capture_video, run_name) for i in range(args.num_envs)]
|
| )
|
|
|
| agent = Agent(envs).to(device)
|
| optimizer = optim.Adam(agent.parameters(), lr=args.learning_rate, eps=1e-5)
|
|
|
| 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
|
| next_obs = torch.Tensor(envs.reset()).to(device)
|
| next_done = torch.zeros(args.num_envs).to(device)
|
| num_updates = args.total_timesteps // args.batch_size
|
|
|
| for update in range(1, num_updates + 1):
|
| for step in range(0, args.num_steps):
|
| 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, done, info = envs.step(action.cpu().numpy())
|
| rewards[step] = torch.tensor(reward).to(device).view(-1)
|
| next_obs, next_done = torch.Tensor(next_obs).to(device), torch.Tensor(done).to(device)
|
|
|
| for item in info:
|
| if "episode" in item.keys():
|
| print(f"global_step={global_step}, episodic_return={item['episode']['r']}")
|
| writer.add_scalar("charts/episodic_return", item["episode"]["r"], global_step)
|
| break
|
|
|
| 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(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,) + envs.single_observation_space.shape)
|
| b_logprobs = logprobs.reshape(-1)
|
| b_actions = actions.reshape((-1,) + envs.single_action_space.shape)
|
| b_advantages = advantages.reshape(-1)
|
| b_returns = returns.reshape(-1)
|
| b_values = values.reshape(-1)
|
|
|
| b_inds = np.arange(args.batch_size)
|
| for epoch in range(args.update_epochs):
|
| np.random.shuffle(b_inds)
|
| for start in range(0, args.batch_size, args.minibatch_size):
|
| end = start + args.minibatch_size
|
| mb_inds = b_inds[start:end]
|
| _, 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_advantages = b_advantages[mb_inds]
|
| 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()
|
|
|
| v_loss = 0.5 * ((newvalue.view(-1) - b_returns[mb_inds]) ** 2).mean()
|
| entropy_loss = entropy.mean()
|
| loss = pg_loss - args.ent_coef * entropy_loss + v_loss * args.vf_coef
|
|
|
| optimizer.zero_grad()
|
| loss.backward()
|
| nn.utils.clip_grad_norm_(agent.parameters(), args.max_grad_norm)
|
| optimizer.step()
|
|
|
| envs.close()
|
| writer.close()
|
|
|
| torch.save(agent.state_dict(), "model.pt")
|
| print("Training completed and model saved as model.pt") |
