比較文本

尋找兩個文字檔案之間的差異

即時編輯器

摺疊未變更行

關閉換行

檢視

比對精度

語法突出顯示

Diffchecker Desktop 執行Diffchecker最安全的方式。取得Diffchecker桌面應用程式：您的差異永遠不會離開您的電腦！取得桌面版

Making PPO work with action masks

建立於 4 年前差異永不過期

1 刪除

行
總計
刪除

字符
總計
刪除

要繼續使用此功能，請升級到 Diffchecker Pro 查看價格

335 行

44 新增

行
總計
新增

字符
總計
新增

要繼續使用此功能，請升級到 Diffchecker Pro 查看價格

366 行

import argparse

import os

import random

import time

from distutils.util import strtobool

import gym

import gym_microrts # fmt: off

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():

# fmt: off

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="MicrortsMining-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=2000000,

help='total timesteps of the experiments')

parser.add_argument('--torch-deterministic', type=lambda x:bool(strtobool(x)), default=True, nargs='?', const=True,

help='if toggled, `torch.backends.cudnn.deterministic=False`')

parser.add_argument('--cuda', type=lambda x:bool(strtobool(x)), default=True, nargs='?', const=True,

help='if toggled, cuda will be enabled by default')

parser.add_argument('--track', type=lambda x:bool(strtobool(x)), default=False, nargs='?', const=True,

help='if toggled, this experiment will be tracked with Weights and Biases')

parser.add_argument('--wandb-project-name', type=str, default="cleanRL",

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, nargs='?', const=True,

help='weather to capture videos of the agent performances (check out `videos` folder)')

# Algorithm specific arguments

parser.add_argument('--num-envs', type=int, default=8,

help='the number of parallel game environments')

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, nargs='?', const=True,

help="Toggle learning rate annealing for policy and value networks")

parser.add_argument('--gae', type=lambda x:bool(strtobool(x)), default=True, nargs='?', const=True,

help='Use GAE for advantage computation')

parser.add_argument('--gamma', type=float, default=0.99,

help='the discount factor gamma')

parser.add_argument('--gae-lambda', type=float, default=0.95,

help='the lambda for the general advantage estimation')

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, nargs='?', const=True,

help="Toggles advantages normalization")

parser.add_argument('--clip-coef', type=float, default=0.1,

help="the surrogate clipping coefficient")

parser.add_argument('--clip-vloss', type=lambda x:bool(strtobool(x)), default=True, nargs='?', const=True,

help='Toggles wheter or not to use a clipped loss for the value function, as per the paper.')

parser.add_argument('--ent-coef', type=float, default=0.01,

help="coefficient of the entropy")

parser.add_argument('--vf-coef', 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)

# fmt: on

return args

def make_env(gym_id, seed, idx, capture_video, run_name):

def thunk():

env = gym.make(gym_id)

env = gym.wrappers.RecordEpisodeStatistics(env)

if capture_video:

if 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 Transpose(nn.Module):

def __init__(self, permutation):

super().__init__()

self.permutation = permutation

def forward(self, x):

return x.permute(self.permutation)

class CategoricalMasked(Categorical):

def __init__(self, probs=None, logits=None, validate_args=None, masks=[]):

self.masks = masks

if len(self.masks) == 0:

super(CategoricalMasked, self).__init__(probs, logits, validate_args)

else:

self.masks = masks.type(torch.BoolTensor).to(device)

logits = torch.where(self.masks, logits, torch.tensor(-1e+8).to(device))

super(CategoricalMasked, self).__init__(probs, logits, validate_args)

def entropy(self):

if len(self.masks) == 0:

return super(CategoricalMasked, self).entropy()

p_log_p = self.logits * self.probs

p_log_p = torch.where(self.masks, p_log_p, torch.tensor(0.).to(device))

return -p_log_p.sum(-1)

class Agent(nn.Module):

def __init__(self, envs):

super(Agent, self).__init__()

self.network = nn.Sequential(

Transpose((0, 3, 1, 2)),

layer_init(nn.Conv2d(27, 16, kernel_size=3, stride=2)),

nn.ReLU(),

layer_init(nn.Conv2d(16, 32, kernel_size=2)),

nn.ReLU(),

nn.Flatten(),

layer_init(nn.Linear(32*3*3, 128)),

nn.ReLU(),

)

self.nvec = envs.single_action_space.nvec

self.actor = layer_init(nn.Linear(128, self.nvec.sum()), std=0.01)

self.critic = layer_init(nn.Linear(128, 1), std=1)

def get_value(self, x):

return self.critic(self.network(x))

def get_action_and_value(self, x, action=None):

def get_action_and_value(self, x, action_mask, action=None):

hidden = self.network(x)

logits = self.actor(hidden)

split_logits = torch.split(logits, self.nvec.tolist(), dim=1)

multi_categoricals = [Categorical(logits=logits) for logits in split_logits]

split_action_masks = torch.split(action_mask, self.nvec.tolist(), dim=1)

multi_categoricals = [

CategoricalMasked(logits=logits, masks=iam)

for (logits, iam) in zip(split_logits, split_action_masks)

]

if action is None:

action = torch.stack([categorical.sample() for categorical in multi_categoricals])

logprob = torch.stack([categorical.log_prob(a) for a, categorical in zip(action, multi_categoricals)])

entropy = torch.stack([categorical.entropy() for categorical in multi_categoricals])

return action.T, logprob.sum(0), entropy.sum(0), self.critic(hidden)

if __name__ == "__main__":

args = parse_args()

run_name = f"{args.gym_id}__{args.exp_name}__{args.seed}__{int(time.time())}"

if args.track:

import wandb

wandb.init(

project=args.wandb_project_name,

entity=args.wandb_entity,

sync_tensorboard=True,

config=vars(args),

name=run_name,

monitor_gym=True,

save_code=True,

)

writer = SummaryWriter(f"runs/{run_name}")

writer.add_text(

"hyperparameters",

"|param|value|\n|-|-|\n%s" % ("\n".join([f"|{key}|{value}|" for key, value in vars(args).items()])),

)

# TRY NOT TO MODIFY: seeding

random.seed(args.seed)

np.random.seed(args.seed)

torch.manual_seed(args.seed)

torch.backends.cudnn.deterministic = args.torch_deterministic

device = torch.device("cuda" if torch.cuda.is_available() and args.cuda else "cpu")

# env setup

envs = gym.vector.SyncVectorEnv(

[make_env(args.gym_id, args.seed + i, i, args.capture_video, run_name) for i in range(args.num_envs)]

)

assert isinstance(envs.single_action_space, gym.spaces.MultiDiscrete), "only MultiDiscrete action space is supported"

agent = Agent(envs).to(device)

optimizer = optim.Adam(agent.parameters(), lr=args.learning_rate, eps=1e-5)

# ALGO Logic: Storage setup

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)

action_masks = torch.zeros((args.num_steps, args.num_envs) + (envs.single_action_space.nvec.sum(),)).to(device)

# TRY NOT TO MODIFY: start the game

global_step = 0

start_time = time.time()

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):

# Annealing the rate if instructed to do so.

if args.anneal_lr:

frac = 1.0 - (update - 1.0) / num_updates

lrnow = frac * args.learning_rate

optimizer.param_groups[0]["lr"] = lrnow

for step in range(0, args.num_steps):

global_step += 1 * args.num_envs

obs[step] = next_obs

dones[step] = next_done

action_masks[step] = torch.Tensor(

np.array([env.action_mask for env in envs.envs])

)

# ALGO LOGIC: action logic

with torch.no_grad():

action, logprob, _, value = agent.get_action_and_value(next_obs)

action, logprob, _, value = agent.get_action_and_value(next_obs, action_masks[step])

values[step] = value.flatten()

actions[step] = action

logprobs[step] = logprob

# TRY NOT TO MODIFY: execute the game and log data.

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)

writer.add_scalar("charts/episodic_length", item["episode"]["l"], global_step)

break

# bootstrap value if not done

with torch.no_grad():

next_value = agent.get_value(next_obs).reshape(1, -1)

if args.gae:

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

else:

returns = torch.zeros_like(rewards).to(device)

for t in reversed(range(args.num_steps)):

if t == args.num_steps - 1:

nextnonterminal = 1.0 - next_done

next_return = next_value

else:

nextnonterminal = 1.0 - dones[t + 1]

next_return = returns[t + 1]

returns[t] = rewards[t] + args.gamma * nextnonterminal * next_return

advantages = returns - 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_advantages = advantages.reshape(-1)

b_returns = returns.reshape(-1)

b_values = values.reshape(-1)

b_action_masks = action_masks.reshape((-1, action_masks.shape[-1]))

# Optimizaing the policy and value network

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):

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].T)

_, newlogprob, entropy, newvalue = agent.get_action_and_value(

b_obs[mb_inds],

b_action_masks[mb_inds],

b_actions.long()[mb_inds].T,

)

logratio = newlogprob - b_logprobs[mb_inds]

ratio = logratio.exp()

with torch.no_grad():

# calculate approx_kl http://joschu.net/blog/kl-approx.html

# old_approx_kl = (-logratio).mean()

approx_kl = ((ratio - 1) - logratio).mean()

clipfracs += [((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)

# Policy loss

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()

# Value loss

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_max = torch.max(v_loss_unclipped, v_loss_clipped)

v_loss = 0.5 * v_loss_max.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 + v_loss * args.vf_coef

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:

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

# TRY NOT TO MODIFY: record rewards for plotting purposes

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("SPS:", int(global_step / (time.time() - start_time)))

writer.add_scalar("charts/SPS", int(global_step / (time.time() - start_time)), global_step)

envs.close()

writer.close()

已保存差異

原始文本

開啟檔案

import argparse
import os
import random
import time
from distutils.util import strtobool

import gym
import gym_microrts # fmt: off
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():
    # fmt: off
    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="MicrortsMining-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=2000000,
        help='total timesteps of the experiments')
    parser.add_argument('--torch-deterministic', type=lambda x:bool(strtobool(x)), default=True, nargs='?', const=True,
        help='if toggled, `torch.backends.cudnn.deterministic=False`')
    parser.add_argument('--cuda', type=lambda x:bool(strtobool(x)), default=True, nargs='?', const=True,
        help='if toggled, cuda will be enabled by default')
    parser.add_argument('--track', type=lambda x:bool(strtobool(x)), default=False, nargs='?', const=True,
        help='if toggled, this experiment will be tracked with Weights and Biases')
    parser.add_argument('--wandb-project-name', type=str, default="cleanRL",
        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, nargs='?', const=True,
        help='weather to capture videos of the agent performances (check out `videos` folder)')

# Algorithm specific arguments
    parser.add_argument('--num-envs', type=int, default=8,
        help='the number of parallel game environments')
    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, nargs='?', const=True,
        help="Toggle learning rate annealing for policy and value networks")
    parser.add_argument('--gae', type=lambda x:bool(strtobool(x)), default=True, nargs='?', const=True,
        help='Use GAE for advantage computation')
    parser.add_argument('--gamma', type=float, default=0.99,
        help='the discount factor gamma')
    parser.add_argument('--gae-lambda', type=float, default=0.95,
        help='the lambda for the general advantage estimation')
    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, nargs='?', const=True,
        help="Toggles advantages normalization")
    parser.add_argument('--clip-coef', type=float, default=0.1,
        help="the surrogate clipping coefficient")
    parser.add_argument('--clip-vloss', type=lambda x:bool(strtobool(x)), default=True, nargs='?', const=True,
        help='Toggles wheter or not to use a clipped loss for the value function, as per the paper.')
    parser.add_argument('--ent-coef', type=float, default=0.01,
        help="coefficient of the entropy")
    parser.add_argument('--vf-coef', 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)
    # fmt: on
    return args

def make_env(gym_id, seed, idx, capture_video, run_name):
    def thunk():
        env = gym.make(gym_id)
        env = gym.wrappers.RecordEpisodeStatistics(env)
        if capture_video:
            if 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 Transpose(nn.Module):
    def __init__(self, permutation):
        super().__init__()
        self.permutation = permutation

def forward(self, x):
        return x.permute(self.permutation)

class Agent(nn.Module):
    def __init__(self, envs):
        super(Agent, self).__init__()
        self.network = nn.Sequential(
            Transpose((0, 3, 1, 2)),
            layer_init(nn.Conv2d(27, 16, kernel_size=3, stride=2)),
            nn.ReLU(),
            layer_init(nn.Conv2d(16, 32, kernel_size=2)),
            nn.ReLU(),
            nn.Flatten(),
            layer_init(nn.Linear(32*3*3, 128)),
            nn.ReLU(),
        )
        self.nvec = envs.single_action_space.nvec
        self.actor = layer_init(nn.Linear(128, self.nvec.sum()), std=0.01)
        self.critic = layer_init(nn.Linear(128, 1), std=1)

def get_value(self, x):
        return self.critic(self.network(x))

def get_action_and_value(self, x, action=None):
        hidden = self.network(x)
        logits = self.actor(hidden)
        split_logits = torch.split(logits, self.nvec.tolist(), dim=1)
        multi_categoricals = [Categorical(logits=logits) for logits in split_logits]
        if action is None:
            action = torch.stack([categorical.sample() for categorical in multi_categoricals])
        logprob = torch.stack([categorical.log_prob(a) for a, categorical in zip(action, multi_categoricals)])
        entropy = torch.stack([categorical.entropy() for categorical in multi_categoricals])
        return action.T, logprob.sum(0), entropy.sum(0), self.critic(hidden)

if __name__ == "__main__":
    args = parse_args()
    run_name = f"{args.gym_id}__{args.exp_name}__{args.seed}__{int(time.time())}"
    if args.track:
        import wandb

wandb.init(
            project=args.wandb_project_name,
            entity=args.wandb_entity,
            sync_tensorboard=True,
            config=vars(args),
            name=run_name,
            monitor_gym=True,
            save_code=True,
        )
    writer = SummaryWriter(f"runs/{run_name}")
    writer.add_text(
        "hyperparameters",
        "|param|value|\n|-|-|\n%s" % ("\n".join([f"|{key}|{value}|" for key, value in vars(args).items()])),
    )

# TRY NOT TO MODIFY: seeding
    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    torch.backends.cudnn.deterministic = args.torch_deterministic

device = torch.device("cuda" if torch.cuda.is_available() and args.cuda else "cpu")

# env setup
    envs = gym.vector.SyncVectorEnv(
        [make_env(args.gym_id, args.seed + i, i, args.capture_video, run_name) for i in range(args.num_envs)]
    )
    assert isinstance(envs.single_action_space, gym.spaces.MultiDiscrete), "only MultiDiscrete action space is supported"

agent = Agent(envs).to(device)
    optimizer = optim.Adam(agent.parameters(), lr=args.learning_rate, eps=1e-5)

# ALGO Logic: Storage setup
    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)

# TRY NOT TO MODIFY: start the game
    global_step = 0
    start_time = time.time()
    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):
        # Annealing the rate if instructed to do so.
        if args.anneal_lr:
            frac = 1.0 - (update - 1.0) / num_updates
            lrnow = frac * args.learning_rate
            optimizer.param_groups[0]["lr"] = lrnow

for step in range(0, args.num_steps):
            global_step += 1 * args.num_envs
            obs[step] = next_obs
            dones[step] = next_done

# ALGO LOGIC: action logic
            with torch.no_grad():
                action, logprob, _, value = agent.get_action_and_value(next_obs)
                values[step] = value.flatten()
            actions[step] = action
            logprobs[step] = logprob

# TRY NOT TO MODIFY: execute the game and log data.
            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)
                    writer.add_scalar("charts/episodic_length", item["episode"]["l"], global_step)
                    break

# bootstrap value if not done
        with torch.no_grad():
            next_value = agent.get_value(next_obs).reshape(1, -1)
            if args.gae:
                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
            else:
                returns = torch.zeros_like(rewards).to(device)
                for t in reversed(range(args.num_steps)):
                    if t == args.num_steps - 1:
                        nextnonterminal = 1.0 - next_done
                        next_return = next_value
                    else:
                        nextnonterminal = 1.0 - dones[t + 1]
                        next_return = returns[t + 1]
                    returns[t] = rewards[t] + args.gamma * nextnonterminal * next_return
                advantages = returns - 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_advantages = advantages.reshape(-1)
        b_returns = returns.reshape(-1)
        b_values = values.reshape(-1)

# Optimizaing the policy and value network
        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):
                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].T)
                logratio = newlogprob - b_logprobs[mb_inds]
                ratio = logratio.exp()

with torch.no_grad():
                    # calculate approx_kl http://joschu.net/blog/kl-approx.html
                    # old_approx_kl = (-logratio).mean()
                    approx_kl = ((ratio - 1) - logratio).mean()
                    clipfracs += [((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)

# Policy loss
                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()

# Value loss
                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_max = torch.max(v_loss_unclipped, v_loss_clipped)
                    v_loss = 0.5 * v_loss_max.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 + v_loss * args.vf_coef

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:
                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

# TRY NOT TO MODIFY: record rewards for plotting purposes
        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("SPS:", int(global_step / (time.time() - start_time)))
        writer.add_scalar("charts/SPS", int(global_step / (time.time() - start_time)), global_step)

envs.close()
    writer.close()

更改後文本

開啟檔案

import argparse
import os
import random
import time
from distutils.util import strtobool

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 Transpose(nn.Module):
    def __init__(self, permutation):
        super().__init__()
        self.permutation = permutation

def forward(self, x):
        return x.permute(self.permutation)

class CategoricalMasked(Categorical):
    def __init__(self, probs=None, logits=None, validate_args=None, masks=[]):
        self.masks = masks
        if len(self.masks) == 0:
            super(CategoricalMasked, self).__init__(probs, logits, validate_args)
        else:
            self.masks = masks.type(torch.BoolTensor).to(device)
            logits = torch.where(self.masks, logits, torch.tensor(-1e+8).to(device))
            super(CategoricalMasked, self).__init__(probs, logits, validate_args)
    
    def entropy(self):
        if len(self.masks) == 0:
            return super(CategoricalMasked, self).entropy()
        p_log_p = self.logits * self.probs
        p_log_p = torch.where(self.masks, p_log_p, torch.tensor(0.).to(device))
        return -p_log_p.sum(-1)

def get_value(self, x):
        return self.critic(self.network(x))

def get_action_and_value(self, x, action_mask, action=None):
        hidden = self.network(x)
        logits = self.actor(hidden)
        split_logits = torch.split(logits, self.nvec.tolist(), dim=1)
        split_action_masks = torch.split(action_mask, self.nvec.tolist(), dim=1)
        multi_categoricals = [
            CategoricalMasked(logits=logits, masks=iam)
            for (logits, iam) in zip(split_logits, split_action_masks)
        ]
        if action is None:
            action = torch.stack([categorical.sample() for categorical in multi_categoricals])
        logprob = torch.stack([categorical.log_prob(a) for a, categorical in zip(action, multi_categoricals)])
        entropy = torch.stack([categorical.entropy() for categorical in multi_categoricals])
        return action.T, logprob.sum(0), entropy.sum(0), self.critic(hidden)

if __name__ == "__main__":
    args = parse_args()
    run_name = f"{args.gym_id}__{args.exp_name}__{args.seed}__{int(time.time())}"
    if args.track:
        import wandb

# TRY NOT TO MODIFY: seeding
    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    torch.backends.cudnn.deterministic = args.torch_deterministic

device = torch.device("cuda" if torch.cuda.is_available() and args.cuda else "cpu")

agent = Agent(envs).to(device)
    optimizer = optim.Adam(agent.parameters(), lr=args.learning_rate, eps=1e-5)

for step in range(0, args.num_steps):
            global_step += 1 * args.num_envs
            obs[step] = next_obs
            dones[step] = next_done
            action_masks[step] = torch.Tensor(
                np.array([env.action_mask for env in envs.envs])
            )

# ALGO LOGIC: action logic
            with torch.no_grad():
                action, logprob, _, value = agent.get_action_and_value(next_obs, action_masks[step])
                values[step] = value.flatten()
            actions[step] = action
            logprobs[step] = logprob

_, newlogprob, entropy, newvalue = agent.get_action_and_value(
                    b_obs[mb_inds],
                    b_action_masks[mb_inds],
                    b_actions.long()[mb_inds].T,
                )
                logratio = newlogprob - b_logprobs[mb_inds]
                ratio = logratio.exp()

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

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()

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

envs.close()
    writer.close()