Making PPO work with MultiDiscrete action spaces

Created Diff never expires
34 removals
338 lines
32 additions
335 lines
import argparse
import argparse
import os
import os
import random
import random
import time
import time
from distutils.util import strtobool
from distutils.util import strtobool


import gym
import gym
import gym_microrts # fmt: off
import numpy as np
import numpy as np
import torch
import torch
import torch.nn as nn
import torch.nn as nn
import torch.optim as optim
import torch.optim as optim
from stable_baselines3.common.atari_wrappers import (
ClipRewardEnv,
EpisodicLifeEnv,
FireResetEnv,
MaxAndSkipEnv,
NoopResetEnv,
)
from torch.distributions.categorical import Categorical
from torch.distributions.categorical import Categorical
from torch.utils.tensorboard import SummaryWriter
from torch.utils.tensorboard import SummaryWriter




def parse_args():
def parse_args():
# fmt: off
# fmt: off
parser = argparse.ArgumentParser()
parser = argparse.ArgumentParser()
parser.add_argument('--exp-name', type=str, default=os.path.basename(__file__).rstrip(".py"),
parser.add_argument('--exp-name', type=str, default=os.path.basename(__file__).rstrip(".py"),
help='the name of this experiment')
help='the name of this experiment')
parser.add_argument('--gym-id', type=str, default="BreakoutNoFrameskip-v4",
parser.add_argument('--gym-id', type=str, default="MicrortsMining-v1",
help='the id of the gym environment')
help='the id of the gym environment')
parser.add_argument('--learning-rate', type=float, default=2.5e-4,
parser.add_argument('--learning-rate', type=float, default=2.5e-4,
help='the learning rate of the optimizer')
help='the learning rate of the optimizer')
parser.add_argument('--seed', type=int, default=1,
parser.add_argument('--seed', type=int, default=1,
help='seed of the experiment')
help='seed of the experiment')
parser.add_argument('--total-timesteps', type=int, default=10000000,
parser.add_argument('--total-timesteps', type=int, default=2000000,
help='total timesteps of the experiments')
help='total timesteps of the experiments')
parser.add_argument('--torch-deterministic', type=lambda x:bool(strtobool(x)), default=True, nargs='?', const=True,
parser.add_argument('--torch-deterministic', type=lambda x:bool(strtobool(x)), default=True, nargs='?', const=True,
help='if toggled, `torch.backends.cudnn.deterministic=False`')
help='if toggled, `torch.backends.cudnn.deterministic=False`')
parser.add_argument('--cuda', type=lambda x:bool(strtobool(x)), default=True, nargs='?', const=True,
parser.add_argument('--cuda', type=lambda x:bool(strtobool(x)), default=True, nargs='?', const=True,
help='if toggled, cuda will be enabled by default')
help='if toggled, cuda will be enabled by default')
parser.add_argument('--track', type=lambda x:bool(strtobool(x)), default=False, nargs='?', const=True,
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')
help='if toggled, this experiment will be tracked with Weights and Biases')
parser.add_argument('--wandb-project-name', type=str, default="cleanRL",
parser.add_argument('--wandb-project-name', type=str, default="cleanRL",
help="the wandb's project name")
help="the wandb's project name")
parser.add_argument('--wandb-entity', type=str, default=None,
parser.add_argument('--wandb-entity', type=str, default=None,
help="the entity (team) of wandb's project")
help="the entity (team) of wandb's project")
parser.add_argument('--capture-video', type=lambda x:bool(strtobool(x)), default=False, nargs='?', const=True,
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)')
help='weather to capture videos of the agent performances (check out `videos` folder)')


# Algorithm specific arguments
# Algorithm specific arguments
parser.add_argument('--num-envs', type=int, default=8,
parser.add_argument('--num-envs', type=int, default=8,
help='the number of parallel game environments')
help='the number of parallel game environments')
parser.add_argument('--num-steps', type=int, default=128,
parser.add_argument('--num-steps', type=int, default=128,
help='the number of steps to run in each environment per policy rollout')
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,
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")
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,
parser.add_argument('--gae', type=lambda x:bool(strtobool(x)), default=True, nargs='?', const=True,
help='Use GAE for advantage computation')
help='Use GAE for advantage computation')
parser.add_argument('--gamma', type=float, default=0.99,
parser.add_argument('--gamma', type=float, default=0.99,
help='the discount factor gamma')
help='the discount factor gamma')
parser.add_argument('--gae-lambda', type=float, default=0.95,
parser.add_argument('--gae-lambda', type=float, default=0.95,
help='the lambda for the general advantage estimation')
help='the lambda for the general advantage estimation')
parser.add_argument('--num-minibatches', type=int, default=4,
parser.add_argument('--num-minibatches', type=int, default=4,
help='the number of mini-batches')
help='the number of mini-batches')
parser.add_argument('--update-epochs', type=int, default=4,
parser.add_argument('--update-epochs', type=int, default=4,
help="the K epochs to update the policy")
help="the K epochs to update the policy")
parser.add_argument('--norm-adv', type=lambda x:bool(strtobool(x)), default=True, nargs='?', const=True,
parser.add_argument('--norm-adv', type=lambda x:bool(strtobool(x)), default=True, nargs='?', const=True,
help="Toggles advantages normalization")
help="Toggles advantages normalization")
parser.add_argument('--clip-coef', type=float, default=0.1,
parser.add_argument('--clip-coef', type=float, default=0.1,
help="the surrogate clipping coefficient")
help="the surrogate clipping coefficient")
parser.add_argument('--clip-vloss', type=lambda x:bool(strtobool(x)), default=True, nargs='?', const=True,
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.')
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,
parser.add_argument('--ent-coef', type=float, default=0.01,
help="coefficient of the entropy")
help="coefficient of the entropy")
parser.add_argument('--vf-coef', type=float, default=0.5,
parser.add_argument('--vf-coef', type=float, default=0.5,
help="coefficient of the value function")
help="coefficient of the value function")
parser.add_argument('--max-grad-norm', type=float, default=0.5,
parser.add_argument('--max-grad-norm', type=float, default=0.5,
help='the maximum norm for the gradient clipping')
help='the maximum norm for the gradient clipping')
parser.add_argument('--target-kl', type=float, default=None,
parser.add_argument('--target-kl', type=float, default=None,
help='the target KL divergence threshold')
help='the target KL divergence threshold')
args = parser.parse_args()
args = parser.parse_args()
args.batch_size = int(args.num_envs * args.num_steps)
args.batch_size = int(args.num_envs * args.num_steps)
args.minibatch_size = int(args.batch_size // args.num_minibatches)
args.minibatch_size = int(args.batch_size // args.num_minibatches)
# fmt: on
# fmt: on
return args
return args




def make_env(gym_id, seed, idx, capture_video, run_name):
def make_env(gym_id, seed, idx, capture_video, run_name):
def thunk():
def thunk():
env = gym.make(gym_id)
env = gym.make(gym_id)
env = gym.wrappers.RecordEpisodeStatistics(env)
env = gym.wrappers.RecordEpisodeStatistics(env)
if capture_video:
if capture_video:
if idx == 0:
if idx == 0:
env = gym.wrappers.RecordVideo(env, f"videos/{run_name}")
env = gym.wrappers.RecordVideo(env, f"videos/{run_name}")
env = NoopResetEnv(env, noop_max=30)
env = MaxAndSkipEnv(env, skip=4)
env = EpisodicLifeEnv(env)
if "FIRE" in env.unwrapped.get_action_meanings():
env = FireResetEnv(env)
env = ClipRewardEnv(env)
env = gym.wrappers.ResizeObservation(env, (84, 84))
env = gym.wrappers.GrayScaleObservation(env)
env = gym.wrappers.FrameStack(env, 4)
env.seed(seed)
env.seed(seed)
env.action_space.seed(seed)
env.action_space.seed(seed)
env.observation_space.seed(seed)
env.observation_space.seed(seed)
return env
return env


return thunk
return thunk




def layer_init(layer, std=np.sqrt(2), bias_const=0.0):
def layer_init(layer, std=np.sqrt(2), bias_const=0.0):
torch.nn.init.orthogonal_(layer.weight, std)
torch.nn.init.orthogonal_(layer.weight, std)
torch.nn.init.constant_(layer.bias, bias_const)
torch.nn.init.constant_(layer.bias, bias_const)
return layer
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):
class Agent(nn.Module):
def __init__(self, envs):
def __init__(self, envs):
super(Agent, self).__init__()
super(Agent, self).__init__()
self.network = nn.Sequential(
self.network = nn.Sequential(
layer_init(nn.Conv2d(4, 32, 8, stride=4)),
Transpose((0, 3, 1, 2)),
nn.ReLU(),
layer_init(nn.Conv2d(27, 16, kernel_size=3, stride=2)),
layer_init(nn.Conv2d(32, 64, 4, stride=2)),
nn.ReLU(),
nn.ReLU(),
layer_init(nn.Conv2d(64, 64, 3, stride=1)),
layer_init(nn.Conv2d(16, 32, kernel_size=2)),
nn.ReLU(),
nn.ReLU(),
nn.Flatten(),
nn.Flatten(),
layer_init(nn.Linear(64 * 7 * 7, 512)),
layer_init(nn.Linear(32*3*3, 128)),
nn.ReLU(),
nn.ReLU(),
)
)
self.actor = layer_init(nn.Linear(512, envs.single_action_space.n), std=0.01)
self.nvec = envs.single_action_space.nvec
self.critic = layer_init(nn.Linear(512, 1), std=1)
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):
def get_value(self, x):
return self.critic(self.network(x / 255.0))
return self.critic(self.network(x))


def get_action_and_value(self, x, action=None):
def get_action_and_value(self, x, action=None):
hidden = self.network(x / 255.0)
hidden = self.network(x)
logits = self.actor(hidden)
logits = self.actor(hidden)
probs = Categorical(logits=logits)
split_logits = torch.split(logits, self.nvec.tolist(), dim=1)
multi_categoricals = [Categorical(logits=logits) for logits in split_logits]
if action is None:
if action is None:
action = probs.sample()
action = torch.stack([categorical.sample() for categorical in multi_categoricals])
return action, probs.log_prob(action), probs.entropy(), self.critic(hidden)
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__":
if __name__ == "__main__":
args = parse_args()
args = parse_args()
run_name = f"{args.gym_id}__{args.exp_name}__{args.seed}__{int(time.time())}"
run_name = f"{args.gym_id}__{args.exp_name}__{args.seed}__{int(time.time())}"
if args.track:
if args.track:
import wandb
import wandb


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


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


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


# env setup
# env setup
envs = gym.vector.SyncVectorEnv(
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)]
[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.Discrete), "only discrete action space is supported"
assert isinstance(envs.single_action_space, gym.spaces.MultiDiscrete), "only MultiDiscrete action space is supported"


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


# ALGO Logic: Storage setup
# ALGO Logic: Storage setup
obs = torch.zeros((args.num_steps, args.num_envs) + envs.single_observation_space.shape).to(device)
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)
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)
logprobs = torch.zeros((args.num_steps, args.num_envs)).to(device)
rewards = 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)
dones = torch.zeros((args.num_steps, args.num_envs)).to(device)
values = 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
# TRY NOT TO MODIFY: start the game
global_step = 0
global_step = 0
start_time = time.time()
start_time = time.time()
next_obs = torch.Tensor(envs.reset()).to(device)
next_obs = torch.Tensor(envs.reset()).to(device)
next_done = torch.zeros(args.num_envs).to(device)
next_done = torch.zeros(args.num_envs).to(device)
num_updates = args.total_timesteps // args.batch_size
num_updates = args.total_timesteps // args.batch_size


for update in range(1, num_updates + 1):
for update in range(1, num_updates + 1):
# Annealing the rate if instructed to do so.
# Annealing the rate if instructed to do so.
if args.anneal_lr:
if args.anneal_lr:
frac = 1.0 - (update - 1.0) / num_updates
frac = 1.0 - (update - 1.0) / num_updates
lrnow = frac * args.learning_rate
lrnow = frac * args.learning_rate
optimizer.param_groups[0]["lr"] = lrnow
optimizer.param_groups[0]["lr"] = lrnow


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


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


# TRY NOT TO MODIFY: execute the game and log data.
# TRY NOT TO MODIFY: execute the game and log data.
next_obs, reward, done, info = envs.step(action.cpu().numpy())
next_obs, reward, done, info = envs.step(action.cpu().numpy())
rewards[step] = torch.tensor(reward).to(device).view(-1)
rewards[step] = torch.tensor(reward).to(device).view(-1)
next_obs, next_done = torch.Tensor(next_obs).to(device), torch.Tensor(done).to(device)
next_obs, next_done = torch.Tensor(next_obs).to(device), torch.Tensor(done).to(device)


for item in info:
for item in info:
if "episode" in item.keys():
if "episode" in item.keys():
print(f"global_step={global_step}, episodic_return={item['episode']['r']}")
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_return", item["episode"]["r"], global_step)
writer.add_scalar("charts/episodic_length", item["episode"]["l"], global_step)
writer.add_scalar("charts/episodic_length", item["episode"]["l"], global_step)
break
break


# bootstrap value if not done
# bootstrap value if not done
with torch.no_grad():
with torch.no_grad():
next_value = agent.get_value(next_obs).reshape(1, -1)
next_value = agent.get_value(next_obs).reshape(1, -1)
if args.gae:
if args.gae:
advantages = torch.zeros_like(rewards).to(device)
advantages = torch.zeros_like(rewards).to(device)
lastgaelam = 0
lastgaelam = 0
for t in reversed(range(args.num_steps)):
for t in reversed(range(args.num_steps)):
if t == args.num_steps - 1:
if t == args.num_steps - 1:
nextnonterminal = 1.0 - next_done
nextnonterminal = 1.0 - next_done
nextvalues = next_value
nextvalues = next_value
else:
else:
nextnonterminal = 1.0 - dones[t + 1]
nextnonterminal = 1.0 - dones[t + 1]
nextvalues = values[t + 1]
nextvalues = values[t + 1]
delta = rewards[t] + args.gamma * nextvalues * nextnonterminal - values[t]
delta = rewards[t] + args.gamma * nextvalues * nextnonterminal - values[t]
advantages[t] = lastgaelam = delta + args.gamma * args.gae_lambda * nextnonterminal * lastgaelam
advantages[t] = lastgaelam = delta + args.gamma * args.gae_lambda * nextnonterminal * lastgaelam
returns = advantages + values
returns = advantages + values
else:
else:
returns = torch.zeros_like(rewards).to(device)
returns = torch.zeros_like(rewards).to(device)
for t in reversed(range(args.num_steps)):
for t in reversed(range(args.num_steps)):
if t == args.num_steps - 1:
if t == args.num_steps - 1:
nextnonterminal = 1.0 - next_done
nextnonterminal = 1.0 - next_done
next_return = next_value
next_return = next_value
else:
else:
nextnonterminal = 1.0 - dones[t + 1]
nextnonterminal = 1.0 - dones[t + 1]
next_return = returns[t + 1]
next_return = returns[t + 1]
returns[t] = rewards[t] + args.gamma * nextnonterminal * next_return
returns[t] = rewards[t] + args.gamma * nextnonterminal * next_return
advantages = returns - values
advantages = returns - values


# flatten the batch
# flatten the batch
b_obs = obs.reshape((-1,) + envs.single_observation_space.shape)
b_obs = obs.reshape((-1,) + envs.single_observation_space.shape)
b_logprobs = logprobs.reshape(-1)
b_logprobs = logprobs.reshape(-1)
b_actions = actions.reshape((-1,) + envs.single_action_space.shape)
b_actions = actions.reshape((-1,) + envs.single_action_space.shape)
b_advantages = advantages.reshape(-1)
b_advantages = advantages.reshape(-1)
b_returns = returns.reshape(-1)
b_returns = returns.reshape(-1)
b_values = values.reshape(-1)
b_values = values.reshape(-1)


# Optimizaing the policy and value network
# Optimizaing the policy and value network
b_inds = np.arange(args.batch_size)
b_inds = np.arange(args.batch_size)
clipfracs = []
clipfracs = []
for epoch in range(args.update_epochs):
for epoch in range(args.update_epochs):
np.random.shuffle(b_inds)
np.random.shuffle(b_inds)
for start in range(0, args.batch_size, args.minibatch_size):
for start in range(0, args.batch_size, args.minibatch_size):
end = start + args.minibatch_size
end = start + args.minibatch_size
mb_inds = b_inds[start:end]
mb_inds = b_inds[start:end]


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


with torch.no_grad():
with torch.no_grad():
# calculate approx_kl http://joschu.net/blog/kl-approx.html
# calculate approx_kl http://joschu.net/blog/kl-approx.html
# old_approx_kl = (-logratio).mean()
# old_approx_kl = (-logratio).mean()
approx_kl = ((ratio - 1) - logratio).mean()
approx_kl = ((ratio - 1) - logratio).mean()
clipfracs += [((ratio - 1.0).abs() > args.clip_coef).float().mean().item()]
clipfracs += [((ratio - 1.0).abs() > args.clip_coef).float().mean().item()]


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


# Policy loss
# Policy loss
pg_loss1 = -mb_advantages * ratio
pg_loss1 = -mb_advantages * ratio
pg_loss2 = -mb_advantages * torch.clamp(ratio, 1 - args.clip_coef, 1 + args.clip_coef)
pg_loss2 = -mb_advantages * torch.clamp(ratio, 1 - args.clip_coef, 1 + args.clip_coef)
pg_loss = torch.max(pg_loss1, pg_loss2).mean()
pg_loss = torch.max(pg_loss1, pg_loss2).mean()


# Value loss
# Value loss
newvalue = newvalue.view(-1)
newvalue = newvalue.view(-1)
if args.clip_vloss:
if args.clip_vloss:
v_loss_unclipped = (newvalue - b_returns[mb_inds]) ** 2
v_loss_unclipped = (newvalue - b_returns[mb_inds]) ** 2
v_clipped = b_values[mb_inds] + torch.clamp(
v_clipped = b_values[mb_inds] + torch.clamp(
newvalue - b_values[mb_inds],
newvalue - b_values[mb_inds],
-args.clip_coef,
-args.clip_coef,
args.clip_coef,
args.clip_coef,
)
)
v_loss_clipped = (v_clipped - b_returns[mb_inds]) ** 2
v_loss_clipped = (v_clipped - b_returns[mb_inds]) ** 2
v_loss_max = torch.max(v_loss_unclipped, v_loss_clipped)
v_loss_max = torch.max(v_loss_unclipped, v_loss_clipped)
v_loss = 0.5 * v_loss_max.mean()
v_loss = 0.5 * v_loss_max.mean()
else:
else:
v_loss = 0.5 * ((newvalue - b_returns[mb_inds]) ** 2).mean()
v_loss = 0.5 * ((newvalue - b_returns[mb_inds]) ** 2).mean()


entropy_loss = entropy.mean()
entropy_loss = entropy.mean()
loss = pg_loss - args.ent_coef * entropy_loss + v_loss * args.vf_coef
loss = pg_loss - args.ent_coef * entropy_loss + v_loss * args.vf_coef


optimizer.zero_grad()
optimizer.zero_grad()
loss.backward()
loss.backward()
nn.utils.clip_grad_norm_(agent.parameters(), args.max_grad_norm)
nn.utils.clip_grad_norm_(agent.parameters(), args.max_grad_norm)
optimizer.step()
optimizer.step()


if args.target_kl is not None:
if args.target_kl is not None:
if approx_kl > args.target_kl:
if approx_kl > args.target_kl:
break
break


y_pred, y_true = b_values.cpu().numpy(), b_returns.cpu().numpy()
y_pred, y_true = b_values.cpu().numpy(), b_returns.cpu().numpy()
var_y = np.var(y_true)
var_y = np.var(y_true)
explained_var = np.nan if var_y == 0 else 1 - np.var(y_true - y_pred) / var_y
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
# 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("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/value_loss", v_loss.item(), global_step)
writer.add_scalar("losses/policy_loss", pg_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/entropy", entropy_loss.item(), global_step)
writer.add_scalar("losses/approx_kl", 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/clipfrac", np.mean(clipfracs), global_step)
writer.add_scalar("losses/explained_variance", explained_var, global_step)
writer.add_scalar("losses/explained_variance", explained_var, global_step)
print("SPS:", int(global_step / (time.time() - start_time)))
print("SPS:", int(global_step / (time.time() - start_time)))
writer.add_scalar("charts/SPS", int(global_step / (time.time() - start_time)), global_step)
writer.add_scalar("charts/SPS", int(global_step / (time.time() - start_time)), global_step)


envs.close()
envs.close()
writer.close()
writer.close()