def train(rank, args, shared_model, optimizer, env_conf):
start_time = time.time()
ptitle('Training Agent: {}'.format(rank))
#log = {}
#setup_logger('{}_train_log'.format(args.env), r'{0}{1}_train_log'.format(
# args.log_dir, args.env))
#log['{}_train_log'.format(args.env)] = logging.getLogger(
# '{}_train_log'.format(args.env))
gpu_id = args.gpu_ids[rank % len(args.gpu_ids)]
torch.manual_seed(args.seed + rank)
if gpu_id >= 0:
torch.cuda.manual_seed(args.seed + rank)
if 'micropolis' in args.env.lower():
env = micropolis_env(args.env, env_conf, args)
if optimizer is None:
if args.optimizer == 'RMSprop':
optimizer = optim.RMSprop(shared_model.parameters(), lr=args.lr)
if args.optimizer == 'Adam':
optimizer = optim.Adam(shared_model.parameters(),
lr=args.lr,
amsgrad=args.amsgrad)
env.seed(args.seed + rank)
player = Agent(None, env, args, None)
player.gpu_id = gpu_id
if 'micropolis' in args.env.lower():
modelInit = getattr(model, args.design_head)
player.model = modelInit(player.env.observation_space.shape[0],
player.env.action_space,
player.env.env.env.MAP_X)
player.lstm_sizes = player.model.getMemorySizes()
else:
player.model = A3Clstm(player.env.observation_space.shape[0],
player.env.action_space)
lstm_size = 512
if 'micropolis' in args.env.lower():
if 'arcade' not in args.env.lower():
lstm_size = (1, 16, env.env.env.MAP_X, env.env.env.MAP_Y)
player.lstm_size = lstm_size
player.state = player.env.reset()
player.state = torch.from_numpy(player.state).float()
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.state = player.state.cuda()
player.model = player.model.cuda()
player.model.train()
player.eps_len += 2
log_counter = 0
while True:
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.model.load_state_dict(shared_model.state_dict())
else:
player.model.load_state_dict(shared_model.state_dict())
num_lstm_layers = len(player.lstm_sizes)
if player.done:
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.cx = [
Variable(torch.zeros(player.lstm_sizes[i]).cuda())
for i in range(num_lstm_layers)
]
player.hx = [
Variable(torch.zeros(player.lstm_sizes[i]).cuda())
for i in range(num_lstm_layers)
]
else:
player.cx = [
Variable(torch.zeros(lstm_sizes[i]))
for i in range(num_lstm_layers)
]
player.hx = [
Variable(torch.zeros(lstm_sizes[i]))
for i in range(num_lstm_layers)
]
else:
player.cx = [
Variable(player.cx[i].data) for i in range(num_lstm_layers)
]
player.hx = [
Variable(player.hx[i].data) for i in range(num_lstm_layers)
]
for step in range(args.num_steps):
player.action_train()
if player.done:
break
if player.done:
state = player.env.reset()
player.state = torch.from_numpy(state).float()
if args.randomize_exploration:
player.certainty = np.random.uniform(0.5, 1.5)
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.state = player.state.cuda()
R = torch.zeros(1, 1)
if not player.done:
values, logit, _ = player.model(
(Variable(player.state.unsqueeze(0)), (player.hx, player.cx)))
if values.size()[1] == 1:
value = values
else:
prob = torch.nn.functional.softmax(logit, dim=1)
action = prob.multinomial(1).data
value = values[0][action]
R = value.data
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
gae = torch.zeros(1, 1).cuda()
R = Variable(R).cuda()
else:
R = Variable(R)
player.values.append(R)
policy_loss = 0
value_loss = 0
for i in reversed(range(len(player.rewards))):
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.rewards[i] = torch.Tensor([player.rewards[i]
]).cuda()
R = args.gamma * R + player.rewards[i]
advantage = R - player.values[i]
value_loss = value_loss + 0.5 * advantage.pow(2)
# Generalized Advantage Estimataion
delta_t = player.rewards[i] + args.gamma * \
player.values[i + 1].data - player.values[i].data
gae = gae * args.gamma * args.tau + delta_t
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
gae = Variable(gae.cuda())
else:
gae = Variable(gae)
policy_loss = policy_loss - \
player.log_probs[i] * Variable(gae) - 0.01 * player.entropies[i]
#if log_counter % 10 == 0:
# log['{}_train_log'.format(args.env)].info(
# "Time {0}, reward {1}, policy loss {2}, value loss {3}, entropy {4}".
# format(time.strftime("%Hh %Mm %Ss", time.gmtime(time.time() - start_time)),
# '{:9.2e}'.format(float(sum(player.rewards) / len(player.rewards))),
# '{:9.2e}'.format(float(policy_loss.data.item())),
# '{:9.2e}'.format(float(value_loss.data.item())),
# '{:10.8e}'.format(float(sum(player.entropies)))))
#log_counter += 1
optimizer.zero_grad()
a3c = args.lmbda * (policy_loss + 0.5 * value_loss)
a3c.backward()
torch.nn.utils.clip_grad_norm_(player.model.parameters(), 40)
ensure_shared_grads(player.model, shared_model, gpu=gpu_id >= 0)
optimizer.step()
player.clear_actions()
def train(rank, args, shared_model, optimizer, env_conf, datasets):
ptitle('Training Agent: {}'.format(rank))
print('Start training agent: ', rank)
if rank == 0:
logger = Logger(args.log_dir)
train_step = 0
gpu_id = args.gpu_ids[rank % len(args.gpu_ids)]
torch.manual_seed(args.seed + rank)
if gpu_id >= 0:
torch.cuda.manual_seed(args.seed + rank)
raw, gt_lbl = datasets
env = EM_env(raw, gt_lbl, env_conf)
if optimizer is None:
if args.optimizer == 'RMSprop':
optimizer = optim.RMSprop(shared_model.parameters(), lr=args.lr)
if args.optimizer == 'Adam':
optimizer = optim.Adam(shared_model.parameters(),
lr=args.lr,
amsgrad=args.amsgrad)
gamma = torch.tensor(1.0)
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
gamma = gamma.cuda()
# env.seed (args.seed + rank)
player = Agent(None, env, args, None)
player.gpu_id = gpu_id
# player.model = A3Clstm (env.observation_space.shape, env_conf["num_action"], args.hidden_feat)
player.model = SimpleCNN(env.observation_space.shape,
env_conf["num_action"])
player.state = player.env.reset()
player.state = torch.from_numpy(player.state).float()
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.state = player.state.cuda()
player.model = player.model.cuda()
player.model.train()
if rank == 0:
eps_reward = 0
pinned_eps_reward = 0
mean_log_prob = 0
while True:
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.model.load_state_dict(shared_model.state_dict())
else:
player.model.load_state_dict(shared_model.state_dict())
if player.done:
player.eps_len = 0
if rank == 0:
pinned_eps_reward = eps_reward
eps_reward = 0
mean_log_prob = 0
for step in range(args.num_steps):
player.action_train()
# print ('step: ', step, 'reward_len: ', len (player.rewards))
if rank == 0:
eps_reward += player.reward
# mean_log_prob += player.log_probs [-1]
if player.done:
break
if player.done:
state = player.env.reset()
player.state = torch.from_numpy(state).float()
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.state = player.state.cuda()
R = torch.zeros(1, 1)
if not player.done:
value, _ = player.model(Variable(player.state.unsqueeze(0)))
R = value.data
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
R = R.cuda()
player.values.append(Variable(R))
policy_loss = 0
value_loss = 0
gae = torch.zeros(1, 1)
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
gae = gae.cuda()
R = Variable(R)
# print ("updating -------------------")
# print ("values:", player.values)
# print ("gamma:", args.gamma)
# print ("rewards:", player.rewards)
for i in reversed(range(len(player.rewards))):
R = args.gamma * R + player.rewards[i]
advantage = R - player.values[i]
value_loss = value_loss + 0.5 * advantage.pow(2)
# print ("advatage: ", advantage)
# print ("value_loss: ", value_loss)
# print ("delta_t: ", player.values[i + 1].data + player.rewards[i])
# Generalized Advantage Estimataion
delta_t = player.values[i + 1].data * args.gamma + player.rewards[i] - \
player.values[i].data
gae = gae * args.gamma * args.tau + delta_t
policy_loss = policy_loss - \
player.log_probs[i] * \
Variable(gae) - 0.01 * player.entropies[i]
player.model.zero_grad()
sum_loss = (policy_loss + value_loss)
sum_loss.backward()
ensure_shared_grads(player.model, shared_model, gpu=gpu_id >= 0)
optimizer.step()
player.clear_actions()
if rank == 0:
train_step += 1
if train_step % (args.log_period) == 0:
log_info = {
'train: sum_loss': sum_loss,
'train: value_loss': value_loss,
'train: policy_loss': policy_loss,
'train: advanage': advantage,
# 'entropy': entropy,
'train: eps reward': pinned_eps_reward,
# 'mean log prob': mean_log_prob
}
for tag, value in log_info.items():
logger.scalar_summary(tag, value, train_step)
def train(rank,
args,
input_model=None,
max_iter=100000,
step_test=-1,
log=False):
if rank >= 0:
torch.manual_seed(args.seed + rank)
gpu_id = args.gpu_ids[rank % len(args.gpu_ids)]
torch.manual_seed(args.seed + rank)
if gpu_id >= 0:
torch.cuda.manual_seed(args.seed + rank)
env = create_env(args)
env.seed(args.seed + rank)
if log:
log = setup_logger("{0}_{1}_log".format(args.scale_legs, rank),
"logs/{0}_{1}_log".format(args.scale_legs, rank))
# player initialization
player = Agent(None, env, args, None)
player.gpu_id = gpu_id
if args.model == 'MLP':
player.model = A3C_MLP(player.env.observation_space.shape[0],
player.env.action_space, args.stack_frames)
if args.model == 'CONV':
player.model = A3C_CONV(args.stack_frames, player.env.action_space)
# load the input model to the player
if input_model != None:
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.model.load_state_dict(input_model.state_dict())
else:
player.model.load_state_dict(input_model.state_dict())
# initialize the player optimizer
optimizer = None
if args.optimizer == 'RMSprop':
optimizer = optim.RMSprop(player.model.dictForOptimizer(), lr=args.lr)
if args.optimizer == 'Adam':
optimizer = optim.Adam(player.model.dictForOptimizer(), lr=args.lr)
else:
optimizer = optim.SGD(player.model.dictForOptimizer(), lr=args.lr)
# reset the environment and initialize the player state
player.state = player.env.reset(args)
player.state = torch.from_numpy(player.state).float()
# If on GPU, do as GPU
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.state = player.state.cuda()
player.model = player.model.cuda()
player.model.train()
last_iter = 0
mean_buf = Buffer(5)
# Start looping over episodes
for iteration in range(max_iter):
last_iter += iteration
# reset cx and hx if the enlvironmnent is over.
if player.done:
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.cx = Variable(torch.zeros(1, 128).cuda())
player.hx = Variable(torch.zeros(1, 128).cuda())
else:
player.cx = Variable(torch.zeros(1, 128))
player.hx = Variable(torch.zeros(1, 128))
else:
player.cx = Variable(player.cx.data)
player.hx = Variable(player.hx.data)
# Roll out actions and collect reward for one episode
for step in range(args.num_steps):
player.action_train()
if player.done:
break
if player.done:
player.eps_len = 0
# reset state
state = player.env.reset(args)
player.state = torch.from_numpy(state).float()
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.state = player.state.cuda()
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
R = torch.zeros(1, 1).cuda()
else:
R = torch.zeros(1, 1)
if not player.done:
state = player.state
if args.model == 'CONV':
state = state.unsqueeze(0)
value, _, _, _ = player.model(
(Variable(state), (player.hx, player.cx)))
R = value.data
player.values.append(Variable(R))
policy_loss = 0
value_loss = 0
R = Variable(R)
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
gae = torch.zeros(1, 1).cuda()
else:
gae = torch.zeros(1, 1)
for i in reversed(range(len(player.rewards))):
R = args.gamma * R + player.rewards[i]
advantage = R - player.values[i]
value_loss = value_loss + 0.5 * advantage.pow(2)
# Generalized Advantage Estimataion
delta_t = player.rewards[i] + args.gamma * \
player.values[i + 1].data - player.values[i].data
gae = gae * args.gamma * args.tau + delta_t
policy_loss = policy_loss - \
(player.log_probs[i].sum() * Variable(gae)) - \
(0.01 * player.entropies[i].sum())
player.model.zero_grad()
(policy_loss + 0.5 * value_loss).backward()
optimizer.step()
player.clear_actions()
if step_test > 0 and iteration % step_test == 0:
tester = Tester(args, player.model)
score = tester.test(last_iter)
mean_buf.push(score)
recent_mean = sum(mean_buf.bf) / mean_buf.current_size
text = "Iteration {0}, episode reward {1}, recent reward mean {2}".format(
iteration, score, recent_mean)
log.info(text)
tester = Tester(args, player.model)
fitness = tester.test(last_iter)
return fitness
def train(rank, args, shared_model, optimizer):
ptitle('Training Agent: {}'.format(rank))
gpu_id = args.gpu_ids[rank % len(args.gpu_ids)]
torch.manual_seed(args.seed + rank)
if gpu_id >= 0:
torch.cuda.manual_seed(args.seed + rank)
env = create_env(args.env, args)
if optimizer is None:
if args.optimizer == 'RMSprop':
optimizer = optim.RMSprop(shared_model.parameters(), lr=args.lr)
if args.optimizer == 'Adam':
optimizer = optim.Adam(shared_model.parameters(), lr=args.lr)
env.seed(args.seed + rank)
player = Agent(None, env, args, None)
player.gpu_id = gpu_id
if args.model == 'MLP':
player.model = A3C_MLP(player.env.observation_space.shape[0],
player.env.action_space, args.stack_frames)
if args.model == 'CONV':
player.model = A3C_CONV(args.stack_frames, player.env.action_space)
player.state = player.env.reset()
player.state = torch.from_numpy(player.state).float()
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.state = player.state.cuda()
player.model = player.model.cuda()
player.model.train()
while True:
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.model.load_state_dict(shared_model.state_dict())
else:
player.model.load_state_dict(shared_model.state_dict())
if player.done:
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.cx = Variable(torch.zeros(1, 128).cuda())
player.hx = Variable(torch.zeros(1, 128).cuda())
else:
player.cx = Variable(torch.zeros(1, 128))
player.hx = Variable(torch.zeros(1, 128))
else:
player.cx = Variable(player.cx.data)
player.hx = Variable(player.hx.data)
for step in range(args.num_steps):
player.action_train()
if player.done:
break
if player.done:
player.eps_len = 0
state = player.env.reset()
player.state = torch.from_numpy(state).float()
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.state = player.state.cuda()
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
R = torch.zeros(1, 1).cuda()
else:
R = torch.zeros(1, 1)
if not player.done:
state = player.state
if args.model == 'CONV':
state = state.unsqueeze(0)
value, _, _, _ = player.model(
(Variable(state), (player.hx, player.cx)))
R = value.data
player.values.append(Variable(R))
policy_loss = 0
value_loss = 0
R = Variable(R)
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
gae = torch.zeros(1, 1).cuda()
else:
gae = torch.zeros(1, 1)
for i in reversed(range(len(player.rewards))):
R = args.gamma * R + player.rewards[i]
advantage = R - player.values[i]
value_loss = value_loss + 0.5 * advantage.pow(2)
# Generalized Advantage Estimataion
# print(player.rewards[i])
delta_t = player.rewards[i] + args.gamma * \
player.values[i + 1].data - player.values[i].data
gae = gae * args.gamma * args.tau + delta_t
policy_loss = policy_loss - \
(player.log_probs[i].sum() * Variable(gae)) - \
(0.01 * player.entropies[i].sum())
player.model.zero_grad()
(policy_loss + 0.5 * value_loss).backward()
ensure_shared_grads(player.model, shared_model, gpu=gpu_id >= 0)
optimizer.step()
player.clear_actions()
def train(rank, args, shared_model, optimizer, env_conf):
ptitle('Training Agent: {}'.format(rank))
print("prank:", rank, "os.pid:", os.getpid())
gpu_id = args.gpu_ids[rank % len(args.gpu_ids)]
torch.manual_seed(args.seed + rank)
if gpu_id >= 0:
torch.cuda.manual_seed(args.seed + rank)
env = AllowBacktracking(
make_local_env(env_conf['game'],
env_conf['level'],
stack=False,
scale_rew=False))
print("Got a local env; obs space:", env.observation_space)
if optimizer is None:
if args.optimizer == 'RMSprop':
optimizer = optim.RMSprop(shared_model.parameters(), lr=args.lr)
if args.optimizer == 'Adam':
optimizer = optim.Adam(shared_model.parameters(),
lr=args.lr,
amsgrad=args.amsgrad)
env.seed(args.seed + rank)
player = Agent(None, env, args, None)
player.gpu_id = gpu_id
player.model = A3Clstm(player.env.observation_space.shape[0],
player.env.action_space)
player.state = player.env.reset()
print("player.state.shape:", player.state.shape)
player.state = torch.from_numpy(player.state).float()
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.state = player.state.cuda()
player.model = player.model.cuda()
player.model.train()
player.eps_len += 2
while True:
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.model.load_state_dict(shared_model.state_dict())
else:
player.model.load_state_dict(shared_model.state_dict())
if player.done:
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.cx = Variable(torch.zeros(1, 512).cuda())
player.hx = Variable(torch.zeros(1, 512).cuda())
else:
player.cx = Variable(torch.zeros(1, 512))
player.hx = Variable(torch.zeros(1, 512))
else:
player.cx = Variable(player.cx.data)
player.hx = Variable(player.hx.data)
for step in range(args.num_steps):
player.action_train()
if player.done:
break
if player.done:
# if player.info['ale.lives'] == 0 or player.max_length:
# player.eps_len = 0
state = player.env.reset()
player.eps_len += 2
player.state = torch.from_numpy(state).float()
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.state = player.state.cuda()
R = torch.zeros(1, 1)
if not player.done:
value, _, _ = player.model(
(Variable(player.state.unsqueeze(0)), (player.hx, player.cx)))
R = value.data
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
R = R.cuda()
player.values.append(Variable(R))
policy_loss = 0
value_loss = 0
gae = torch.zeros(1, 1)
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
gae = gae.cuda()
R = Variable(R)
for i in reversed(range(len(player.rewards))):
R = args.gamma * R + player.rewards[i]
advantage = R - player.values[i]
value_loss = value_loss + 0.5 * advantage.pow(2)
# Generalized Advantage Estimataion
delta_t = player.rewards[i] + args.gamma * \
player.values[i + 1].data - player.values[i].data
gae = gae * args.gamma * args.tau + delta_t
policy_loss = policy_loss - \
player.log_probs[i] * \
Variable(gae) - 0.01 * player.entropies[i]
player.model.zero_grad()
(policy_loss + 0.5 * value_loss).backward()
torch.nn.utils.clip_grad_norm(player.model.parameters(), 100.0)
ensure_shared_grads(player.model, shared_model, gpu=gpu_id >= 0)
optimizer.step()
player.clear_actions()
def train(args, envs, observation_space, action_space):
gpu_id = 0
#每个单独的work,独立的环境和model，在cuda中运行
player = Agent(envs, args)
player.model = A3Clstm(observation_space, action_space)
player.state = player.env.reset()
player.state = torch.from_numpy(player.state).float()
with torch.cuda.device(gpu_id):
player.model = player.model.cuda()
player.state = player.state.cuda()
player.cx = torch.zeros(args.workers, 512).cuda()
player.hx = torch.zeros(args.workers, 512).cuda()
optimizer = torch.optim.Adam(player.model.parameters(),
lr=args.lr,
amsgrad=args.amsgrad)
#切换到训练模式
player.model.train()
while True:
#训练20步或者game over就结束训练
for step in range(args.num_steps):
#训练时，保存每一步的相关信息到list
player.env.get_images()
player.action_train()
if player.dones[-1][0]:
break
if not player.dones[-1][0]:
value, _, _ = player.model((player.state, (player.hx, player.cx)))
R = value.detach()
else:
R = torch.zeros(args.workers, 1)
with torch.cuda.device(gpu_id):
R = R.cuda()
player.values.append(R)
for j in range(args.num_ppo_train):
policy_loss = 0
value_loss = 0
gae = 0
for i in reversed(range(len(player.rewards))):
value, logit, _ = player.model(
(player.states[i], (player.hxs[i], player.cxs[i])))
prob = F.softmax(logit, dim=1)
log_prob = F.log_softmax(logit, dim=1)
entropy = -(log_prob * prob).sum(1)
log_probs_current = log_prob.gather(1, player.actions[i])
R = args.gamma * R + player.rewards[i]
advantage = R - value
value_loss = value_loss + 0.5 * advantage.pow(2)
# Generalized Advantage Estimataion
delta_t = player.rewards[i] + args.gamma * player.values[
i + 1].detach() - player.values[i].detach()
gae = gae * args.gamma * args.tau + delta_t
ratio = torch.exp(log_probs_current - player.log_probs[i])
surr1 = ratio
surr2 = torch.clamp(ratio, 1.0 - args.clip_param,
1.0 + args.clip_param)
policy_loss = policy_loss - torch.min(
surr1, surr2) * gae * -0.01 * entropy
optimizer.zero_grad()
(policy_loss + 0.5 * value_loss).mean().backward()
optimizer.step()
#game over时记忆值重置
if player.dones[-1][0]:
with torch.cuda.device(gpu_id):
player.cx = torch.zeros(args.workers, 512).cuda()
player.hx = torch.zeros(args.workers, 512).cuda()
else:
player.cx = player.cx.detach()
player.hx = player.hx.detach()
player.clear_actions()
# advantage[0:n]
# 第0，1，2，...n 到 n+1的估值差 r[0-n],r[1-n],r[2-n]....rn Value(N+1) 取反：
# 第n,n-1,n-2,n-3,......3,2,1
# r[n] + Value(N+1) - Value(N)
# r[n:n-1] + Value(N+1) - Value(N-1)
# ...
# r[n:2] + Value(N + 1) - Value(2)
# r[n:1] + Value(N + 1) - Value(1)
# R = args.gamma * R + player.rewards[i]
# advantage = R - player.values[i]
# value_loss = value_loss + 0.5 * advantage.pow(2)
# value_loss = 0.5 * advantage.pow(2)
# advantage = args.gamma * R + player.rewards[i] - player.values[i]
#entropy = -(log_prob * prob).sum(1)
#self.entropies.append(entropy)
#通过prob 采样对应的动作和动作logprob
# 计算每次的概率和entropy(entropies)和entropy的sum,sum是每一步所有动作概率的熵值
def train(rank, args, shared_model, optimizer, train_modes, n_iters, env=None):
n_steps = 0
n_iter = 0
writer = SummaryWriter(os.path.join(args.log_dir, 'Agent:{}'.format(rank)))
ptitle('Training Agent: {}'.format(rank))
gpu_id = args.gpu_ids[rank % len(args.gpu_ids)]
torch.manual_seed(args.seed + rank)
training_mode = args.train_mode
env_name = args.env
train_modes.append(training_mode)
n_iters.append(n_iter)
if gpu_id >= 0:
torch.cuda.manual_seed(args.seed + rank)
device = torch.device('cuda:' + str(gpu_id))
else:
device = torch.device('cpu')
if env == None:
env = create_env(env_name)
params = shared_model.parameters()
if optimizer is None:
if args.optimizer == 'RMSprop':
optimizer = optim.RMSprop(params, lr=args.lr)
if args.optimizer == 'Adam':
optimizer = optim.Adam(params, lr=args.lr)
env.seed(args.seed + rank)
player = Agent(None, env, args, None, device)
player.gpu_id = gpu_id
player.env.reset()
# prepare model
player.model = build_model(action_space=player.env.action_space,
pose_space=player.reset_cam_pose(),
args=args,)
player.model = player.model.to(device)
player.model.train()
player.reset()
reward_sum = torch.zeros(player.num_agents).to(device)
count_eps = 0
print('Start training...')
while True:
# sys to the shared model
player.model.load_state_dict(shared_model.state_dict())
if player.done:
player.reset()
reward_sum = torch.zeros(player.num_agents).to(device)
count_eps += 1
player.update_rnn_hidden()
fps_counter = 0
t0 = time.time()
for i in range(args.num_steps):
player.action_train()
reward_sum += player.reward
fps_counter += 1
n_steps += 1
if player.done:
for i, r_i in enumerate(reward_sum):
# add for Pose Only
if i not in player.env.random_ids:
continue
#
writer.add_scalar('train/reward_' + str(i), r_i, n_steps)
break
fps = fps_counter / (time.time() - t0)
policy_loss, value_loss, entropies, pred_loss, values0 = player.optimize(params, optimizer, shared_model, gpu_id)
writer.add_scalar('train/policy_loss_sum', policy_loss.sum(), n_steps)
writer.add_scalar('train/value_loss_sum', value_loss.sum(), n_steps)
writer.add_scalar('train/entropies_sum', entropies.sum(), n_steps)
writer.add_scalar('train/values0', values0.sum(), n_steps)
writer.add_scalar('train/pred_R_loss', pred_loss, n_steps)
writer.add_scalar('train/fps', fps, n_steps)
# writer.add_scalar('train/lr', lr[0], n_iter)
n_iter += 1
n_iters[rank] = n_iter
if train_modes[rank] == -100:
env.close()
break
def train(rank, args, shared_model, optimizer, train_modes, n_iters, device, env=None):
n_steps = 0
n_iter = 0
writer = SummaryWriter(os.path.join(args.log_dir, 'Agent:{}'.format(rank)))
ptitle('Training Agent: {}'.format(rank))
torch.manual_seed(args.seed + rank)
training_mode = args.train_mode
env_name = args.env
train_modes.append(training_mode)
n_iters.append(n_iter)
if env == None:
env = create_env(env_name, args)
params = shared_model.parameters()
if optimizer is None:
if args.optimizer == 'RMSprop':
optimizer = optim.RMSprop(params, lr=args.lr)
if args.optimizer == 'Adam':
optimizer = optim.Adam(filter(lambda p: p.requires_grad, shared_model.parameters()), lr=args.lr)
env.seed(args.seed + rank)
player = Agent(None, env, args, None, None, device)
player.model = build_model(
player.env.observation_space, player.env.action_space, args, device).to(device)
player.state = player.env.reset()
if 'Unreal' in args.env:
player.cam_pos = player.env.env.env.env.cam_pose
player.collect_state = player.env.env.env.env.current_states
player.set_cam_info()
player.state = torch.from_numpy(player.state).float()
player.state = player.state.to(device)
player.model = player.model.to(device)
player.model.train()
reward_sum = torch.zeros(player.num_agents).to(device)
count_eps = 0
cross_entropy_loss = nn.CrossEntropyLoss()
while True:
player.model.load_state_dict(shared_model.state_dict())
player.update_lstm()
fps_counter = 0
t0 = time.time()
for step in range(args.num_steps):
player.action_train()
n_steps += 1
reward_sum += player.reward
if player.done:
break
update_steps = len(player.rewards)
fps = fps_counter / (time.time() - t0)
if player.done:
for i in range(player.num_agents):
writer.add_scalar('train/reward_'+str(i), reward_sum[i], n_steps)
count_eps += 1
reward_sum = torch.zeros(player.num_agents).to(device)
player.eps_len = 0
player.state = player.env.reset()
player.set_cam_info()
player.state = torch.from_numpy(player.state).float().to(device)
R = torch.zeros(player.num_agents, 1).to(device)
if not player.done:
state = player.state
value_multi, _, _, _, _, _, _, _ , _= player.model(
(Variable(state, requires_grad=True),
Variable((player.cam_info), requires_grad=True), player.H_multi,
player.last_gate_ids, player.gt_gate))
for i in range(player.num_agents):
R[i][0] = value_multi[i].data
gates, gt_gates = [], []
for k1 in range(len(player.rewards)):
for k2 in range(player.num_agents):
gates.append(player.gates[k1][k2])
gt_gates.append(player.gate_gts[k1][k2])
gate_probs = torch.cat(gates).view(-1, 2).to(device)
gate_gt_ids = torch.Tensor(gt_gates).view(1, -1).squeeze().long().to(device)
gate_loss = cross_entropy_loss(gate_probs, gate_gt_ids)
player.values.append(Variable(R).to(device))
policy_loss = torch.zeros(player.num_agents, 1).to(device)
value_loss = torch.zeros(player.num_agents, 1).to(device)
entropies = torch.zeros(player.num_agents, 1).to(device)
w_entropies = torch.Tensor([[float(args.entropy)] for i in range(player.num_agents)]).to(device)
R = Variable(R, requires_grad=True).to(device)
gae = torch.zeros(1, 1).to(device)
for i in reversed(range(len(player.rewards))):
R = args.gamma * R + player.rewards[i]
advantage = R - player.values[i]
value_loss = value_loss + 0.5 * advantage.pow(2)
# Generalized Advantage Estimataion
delta_t = player.rewards[i] + args.gamma * player.values[i + 1].data - player.values[i].data
gae = gae * args.gamma * args.tau + delta_t
policy_loss = policy_loss - \
(player.log_probs[i] * Variable(gae)) - \
(w_entropies * player.entropies[i])
entropies += player.entropies[i]
loss = policy_loss.sum() / update_steps / player.num_agents + 0.5 * value_loss.sum() / update_steps / player.num_agents + \
5 * gate_loss
player.model.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm_(params, 50)
ensure_shared_grads(player.model, shared_model, gpu=args.gpu_ids[-1] >= 0)
writer.add_scalar('train/policy_loss_sum', policy_loss.sum(), n_steps)
writer.add_scalar('train/value_loss_sum', value_loss.sum(), n_steps)
writer.add_scalar('train/entropies_sum', entropies.sum(), n_steps)
writer.add_scalar('train/fps', fps, n_steps)
writer.add_scalar('train/gate_loss', gate_loss, n_steps)
n_iter += 1
n_iters[rank] = n_iter
optimizer.step()
player.clear_actions()
if train_modes[rank] == -100:
env.close()
break
def train(rank, args, shared_model, optimizer, env_conf):
ptitle('Train {0}: {1}'.format(args.env, rank))
print('Start training agent: ', rank)
if rank == 0:
logger = Logger(args.log_dir + '_losses/')
train_step = 0
gpu_id = args.gpu_ids[rank % len(args.gpu_ids)]
env_conf["env_gpu"] = gpu_id
torch.manual_seed(args.seed + rank)
if gpu_id >= 0:
torch.cuda.manual_seed(args.seed + rank)
env = database_env(env_conf, seed=0)
if optimizer is None:
if args.optimizer == 'RMSprop':
optimizer = optim.RMSprop(shared_model.parameters(), lr=args.lr)
if args.optimizer == 'Adam':
optimizer = optim.Adam(shared_model.parameters(),
lr=args.lr,
amsgrad=args.amsgrad)
player = Agent(None, env, args, None, gpu_id)
player.gpu_id = gpu_id
player.model = get_model(args,
args.model,
env_conf["observation_shape"],
args.features,
env_conf["num_actions"],
gpu_id=0,
lstm_feats=args.lstm_feats)
player.state = player.env.reset()
player.state = torch.from_numpy(player.state).float()
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.state = player.state.cuda()
player.model = player.model.cuda()
player.model.train()
if rank == 0:
eps_reward = 0
pinned_eps_reward = 0
while True:
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.model.load_state_dict(shared_model.state_dict())
else:
player.model.load_state_dict(shared_model.state_dict())
if player.done:
player.eps_len = 0
if rank == 0:
if train_step % args.train_log_period == 0 and train_step > 0:
print("train: step", train_step, "\teps_reward",
eps_reward)
if train_step > 0:
pinned_eps_reward = player.env.sum_reward
eps_reward = 0
if args.lstm_feats:
player.cx, player.hx = init_linear_lstm(
args.lstm_feats, gpu_id)
elif args.lstm_feats:
player.cx = Variable(player.cx.data)
player.hx = Variable(player.hx.data)
for step in range(args.num_steps):
player.action_train()
if rank == 0:
eps_reward = player.env.sum_reward
if player.done:
break
if player.done:
if rank == 0:
if train_step % args.train_log_period == 0 and train_step > 0:
print("train: step", train_step, "\teps_reward",
eps_reward)
# print ("rewards: ", player.env.rewards)
# print ("actions: ", player.actions)
if player.done:
state = player.env.reset()
player.state = torch.from_numpy(state).float()
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.state = player.state.cuda()
R = torch.zeros(1, 1, 1, 1)
if not player.done:
if args.lstm_feats:
value, _, _ = player.model(
(Variable(player.state.unsqueeze(0)), (player.hx,
player.cx)))
else:
value, _ = player.model(Variable(player.state.unsqueeze(0)))
R = value.data
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
R = R.cuda()
player.values.append(Variable(R))
policy_loss = 0
value_loss = 0
gae = torch.zeros(1, 1, 1, 1)
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
gae = gae.cuda()
R = Variable(R)
for i in reversed(range(len(player.rewards))):
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
reward_i = torch.tensor(player.rewards[i]).cuda()
else:
reward_i = torch.tensor(player.rewards[i])
R = args.gamma * R + reward_i
advantage = R - player.values[i]
value_loss = value_loss + (0.5 * advantage * advantage).mean()
delta_t = player.values[
i + 1].data * args.gamma + reward_i - player.values[i].data
gae = gae * args.gamma * args.tau + delta_t
policy_loss = policy_loss - \
(player.log_probs[i] * Variable(gae)).mean () - \
(args.entropy_alpha * player.entropies[i]).mean ()
player.model.zero_grad()
sum_loss = (policy_loss + value_loss)
sum_loss.backward()
ensure_shared_grads(player.model, shared_model, gpu=gpu_id >= 0)
optimizer.step()
player.clear_actions()
if rank == 0:
train_step += 1
if train_step % args.log_period == 0 and train_step > 0:
log_info = {
'sum_loss': sum_loss,
'value_loss': value_loss,
'policy_loss': policy_loss,
'advanage': advantage,
'train eps reward': pinned_eps_reward,
}
for tag, value in log_info.items():
logger.scalar_summary(tag, value, train_step)
def train(rank, args, shared_model, optimizer):
ptitle('Training Agent: {}'.format(rank))
gpu_id = args.gpu_ids[rank % len(args.gpu_ids)]
writer = SummaryWriter(log_dir=args.log_dir + 'tb_train')
log = {}
setup_logger('{}_train_log'.format(rank),
r'{0}{1}_train_log'.format(args.log_dir, rank))
log['{}_train_log'.format(rank)] = logging.getLogger(
'{}_train_log'.format(rank))
torch.manual_seed(args.seed + rank)
if gpu_id >= 0:
torch.cuda.manual_seed(args.seed + rank)
env = atari_env(env_id=rank, args=args, type='train')
if optimizer is None:
if args.optimizer == 'RMSprop':
optimizer = optim.RMSprop(shared_model.parameters(), lr=args.lr)
if args.optimizer == 'Adam':
optimizer = optim.Adam(shared_model.parameters(),
lr=args.lr,
amsgrad=args.amsgrad)
env.seed(args.seed + rank)
player = Agent(None, env, args, None)
player.gpu_id = gpu_id
player.model = A3Clstm(player.env.observation_space.shape[2],
player.env.action_space.n)
player.state = player.env.reset()
player.state = normalize_rgb_obs(player.state)
player.state = torch.from_numpy(player.state).float()
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.state = player.state.cuda()
player.model = player.model.cuda()
player.model.train()
num_trains = 0
if not os.path.exists(args.log_dir + "images/"):
os.makedirs(args.log_dir + "images/")
while True:
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.model.load_state_dict(shared_model.state_dict())
else:
player.model.load_state_dict(shared_model.state_dict())
for step in range(args.num_steps):
player.action_train()
if player.done:
break
if player.done:
num_trains += 1
log['{}_train_log'.format(rank)].info('entropy:{0}'.format(
player.entropy.data[0]))
writer.add_scalar("data/entropy_" + str(rank),
player.entropy.data[0], num_trains)
writer.add_image('FCN_' + str(rank), player.fcn, num_trains)
writer.add_image('Depth_GroundTruth_' + str(rank), player.depth,
num_trains)
writer.add_image('RGB_' + str(rank), player.env.get_rgb(),
num_trains)
save_image(
player.fcn.data, args.log_dir + "images/" + str(rank) + "_" +
str(num_trains) + "_fcn.png")
# print("player.fcn.data:", player.fcn.data)
save_image(
player.depth.data, args.log_dir + "images/" + str(rank) + "_" +
str(num_trains) + "_depth.png")
cv2.imwrite(
args.log_dir + "images/" + str(rank) + "_" + str(num_trains) +
"_rgb.png", player.env.get_rgb())
# print("player.depth.data:", player.depth.data)
player.eps_len = 0
player.current_life = 0
state = player.env.reset()
state = normalize_rgb_obs(state)
player.state = torch.from_numpy(state).float()
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.state = player.state.cuda()
R = torch.zeros(1, 1)
if not player.done:
with torch.cuda.device(gpu_id):
value, _, _, _ = player.model(
(Variable(player.state.unsqueeze(0)), (player.hx,
player.cx),
Variable(
torch.from_numpy(player.env.target).type(
torch.FloatTensor).cuda())))
R = value.data
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
R = R.cuda()
player.values.append(Variable(R))
policy_loss = 0
value_loss = 0
gae = torch.zeros(1, 1)
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
gae = gae.cuda()
R = Variable(R)
for i in reversed(range(len(player.rewards))):
R = args.gamma * R + player.rewards[i]
advantage = R - player.values[i]
value_loss = value_loss + 0.5 * advantage.pow(2)
# Generalized Advantage Estimataion
delta_t = args.gamma * player.values[
i + 1].data + player.rewards[i] - player.values[i].data
gae = gae * args.gamma * args.tau + delta_t
# policy_loss = policy_loss - \
# player.log_probs[i] * \
# Variable(gae) - 0.01 * player.entropies[i] \
# + player.fcn_losses[i] # FCN
policy_loss = policy_loss - 1e-5*(player.log_probs[i] * Variable(gae)) - 1e-5*(0.01 * player.entropies[i]) \
+ player.fcn_losses[i] * DEPTH_LOSS_DISCOUNT # FCN
# policy_loss = policy_loss + player.fcn_losses[i] # FCN
writer.add_scalar("data/value_loss_" + str(rank), value_loss,
num_trains)
writer.add_scalar("data/policy_loss_" + str(rank), policy_loss,
num_trains)
player.model.zero_grad()
(policy_loss + 0.5 * value_loss).backward()
torch.nn.utils.clip_grad_norm(player.model.parameters(), 40.0)
ensure_shared_grads(player.model, shared_model, gpu=gpu_id >= 0)
optimizer.step()
player.clear_actions()
def train(rank, args, shared_model, optimizer): # optimizer为shared_model的
init = True
ptitle('Training Agent: {}'.format(rank))
torch.manual_seed(args.seed + rank)
env = create_env(args.env, args.seed + rank)
# env = gym.make(args.env)
# env.seed(args.seed + rank)
if optimizer is None:
if args.optimizer == 'RMSprop':
optimizer = optim.RMSprop(shared_model.parameters(), lr=args.lr)
if args.optimizer == 'Adam':
optimizer = optim.Adam(shared_model.parameters(), lr=args.lr)
player = Agent(None, env, args, None, rank)
player.model = A3C_MLP(player.env.observation_space,
player.env.action_space, args.stack_frames)
player.state = player.env.reset()
player.state = torch.from_numpy(player.state).float()
player.model.train() #固定使用场景为train
if rank == 1:
# file = open(os.path.join(args.log_dir, 'TD_Error.txt'), 'w+')
writer = SummaryWriter('8_27_train')
local_step_counter = 0
while True:
if init: # 初始化
shared_model.training_steps.weight.data \
.copy_(torch.Tensor([0]))
shared_model.training_steps.bias.data \
.copy_(torch.Tensor([0]))
init = False
player.model.load_state_dict(
shared_model.state_dict()) # synchronize parameters
for step in range(args.num_steps):
# print("thread", rank, local_step_counter, shared_model.training_steps.weight.data.cpu().numpy())
local_step_counter += 1 # update step counters
shared_model.training_steps.weight.data \
.copy_(torch.Tensor([1]) + shared_model.training_steps.weight.data) # 总步骤（各个worker所走步数之和）T每次加一
player.action_train() # core
if player.done:
break
terminal = False
if player.done or player.eps_len >= args.max_episode_length: # 玩家完成或者超出最大迭代次数
terminal = True
shared_model.done_nums += 1
if 'is_success' in player.info.keys():
shared_model.success_num += 1
R = torch.zeros(1)
if not player.done: # 结算
state = player.state
# A3C，value和policy net是用的同一个网络
value, _, _ = player.model(Variable(state))
R = value.data
if terminal: #重置
shared_model.training_steps.bias.data \
.copy_(torch.Tensor([1]) + shared_model.training_steps.bias.data) # 总步数加一
player.eps_len = 0
state = player.env.reset()
player.state = torch.from_numpy(state).float()
player.reset_flag = True
player.values.append(Variable(R))
policy_loss = 0
value_loss = 0
R = Variable(R)
gae = torch.zeros(1, 1)
for i in reversed(range(len(player.rewards))):
R = args.gamma * R + np.float(player.rewards[i]) # reward
advantage = R - player.values[i] # advantage
value_loss = value_loss + 0.5 * advantage.pow(2) # 公式(10) 更新w
if rank == 1:
# file.write(str(advantage.pow(2).data.cpu().numpy()[0]))
# file.write(' ')
# file.write(
# str(int(shared_model.training_steps.weight.data.cpu().numpy()[0])))
# file.write('\n')
writer.add_scalar(
'TD-error/train',
advantage.pow(2).data.cpu().numpy()[0],
shared_model.training_steps.weight.data.cpu().numpy()[0])
player.values[i] = player.values[i].float()
player.values[i + 1] = player.values[i + 1].float()
delta_t = player.rewards[i] + args.gamma * \
player.values[i + 1].data - \
player.values[i].data # a2c计算td-error
# GAE算法
gae = gae * args.gamma * args.tau + delta_t
policy_loss = policy_loss - \
(player.log_probs[i].sum() * Variable(gae)) - \
(0.01 * player.entropies[i].sum()) # 更新theta 公式(9)
"""
每个线程和环境交互到一定量的数据后，就计算在自己线程里的神经网络损失函数的梯度，
但是这些梯度却并不更新自己线程里的神经网络，而是去更新公共的神经网络。
也就是n个线程会独立的使用累积的梯度分别更新公共部分的神经网络模型参数。
每隔一段时间，线程会将自己的神经网络的参数更新为公共神经网络的参数，进而指导后面的环境交互。
"""
player.model.zero_grad()
# policy_loss + 0.5 * value_loss即为loss
if rank == 1:
writer.add_scalar(
'VLoss/train', value_loss,
shared_model.training_steps.weight.data.cpu().numpy()[0])
writer.add_scalar(
'PLoss/train', policy_loss,
shared_model.training_steps.weight.data.cpu().numpy()[0])
(policy_loss + 0.5 * value_loss).backward() # 计算该worder的损失函数梯度
ensure_shared_grads(player.model, shared_model) # 该worker将自己的参数传给公用的模型
optimizer.step(
) # optimizer为shared_model的 step()将参数更新值施加到shared_model的parameters 上
player.clear_actions()
if shared_model.training_steps.weight.data.cpu().numpy(
) > args.training_steps:
print('num of success={0},training episodes={1},success_rate={2}'.
format(shared_model.success_num, shared_model.done_nums,
shared_model.success_num / shared_model.done_nums))
break
def test(args, shared_model, optimizer, env_conf):
ptitle('Test Agent')
gpu_id = args.gpu_ids[-1]
start_time = datetime.now().strftime('%Y-%m-%d_%H_%M_%S')
log = {}
setup_logger('{}_log'.format(args.env), r'{0}{1}_{2}_log'.format(
args.log_dir, args.env, start_time))
log['{}_log'.format(args.env)] = logging.getLogger('{}_log'.format(
args.env))
d_args = vars(args)
for k in d_args.keys():
log['{}_log'.format(args.env)].info('{0}: {1}'.format(k, d_args[k]))
if not os.path.exists(args.save_model_dir):
os.mkdir(args.save_model_dir)
if args.seed:
torch.manual_seed(args.seed)
if gpu_id >= 0:
torch.cuda.manual_seed(args.seed)
env = atari_env(args.env, env_conf, args)
reward_sum = 0
start = time.time()
num_tests = 0
reward_total_sum = 0
player = Agent(None, env, args, None)
player.gpu_id = gpu_id
player.model = A3Cff(player.env.observation_space.shape[0],
player.env.action_space)
player.state = player.env.reset()
player.eps_len += 2
player.state = torch.from_numpy(player.state).float()
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.model = player.model.cuda()
player.state = player.state.cuda()
flag = True
max_score = -10000
while True:
p = optimizer.param_groups[0]['params'][0]
step = optimizer.state[p]['step']
player.model.eval()
if flag:
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.model.load_state_dict(shared_model.state_dict())
else:
player.model.load_state_dict(shared_model.state_dict())
flag = False
with torch.no_grad():
if args.robust:
#player.action_test_losses(args.epsilon_end)
lin_coeff = min(1, (1.5*int(step)+1)/(args.total_frames/args.num_steps))
epsilon = lin_coeff*args.epsilon_end
player.action_train(epsilon)
else:
player.action_train()
#player.action_test_losses()
reward_sum += player.noclip_reward
if player.done and not player.info:
state = player.env.reset()
player.eps_len += 2
player.state = torch.from_numpy(state).float()
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.state = player.state.cuda()
elif player.info:
# calculate losses for tracking
R = torch.zeros(1, 1)
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
R = R.cuda()
player.values.append(R)
gae = torch.zeros(1, 1)
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
gae = gae.cuda()
R = Variable(R)
standard_loss = 0
worst_case_loss = 0
value_loss = 0
entropy = 0
for i in reversed(range(len(player.rewards))):
R = args.gamma * R + player.rewards[i]
advantage = R - player.values[i]
value_loss += 0.5 * advantage.pow(2)
# Generalized Advantage Estimataion
delta_t = player.rewards[i] + args.gamma * \
player.values[i + 1].data - player.values[i].data
gae = gae * args.gamma * args.tau + delta_t
if args.robust:
if advantage >= 0:
worst_case_loss += - player.min_log_probs[i] * Variable(gae)
else:
worst_case_loss += - player.max_log_probs[i] * Variable(gae)
standard_loss += -player.log_probs[i] * Variable(gae)
entropy += player.entropies[i]
standard_loss = standard_loss/len(player.rewards)
worst_case_loss = worst_case_loss/len(player.rewards)
value_loss = value_loss/len(player.rewards)
entropy = entropy/len(player.rewards)
player.clear_actions()
flag = True
num_tests += 1
reward_total_sum += reward_sum
reward_mean = reward_total_sum / num_tests
log['{}_log'.format(args.env)].info(
("Time {0}, steps {1}/{2}, ep reward {3}, ep length {4}, reward mean {5:.3f} \n"+
"Losses: Policy:{6:.3f}, Worst case: {7:.3f}, Value: {8:.3f}, Entropy: {9:.3f}").
format(time.strftime("%Hh %Mm %Ss", time.gmtime(time.time() - start)),
int(step), args.total_frames/args.num_steps, reward_sum, player.eps_len, reward_mean,
float(standard_loss), float(worst_case_loss), float(value_loss), float(entropy)))
if args.save_max and reward_sum >= max_score:
max_score = reward_sum
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
state_to_save = player.model.state_dict()
torch.save(state_to_save, '{0}{1}_{2}_best.pt'.format(
args.save_model_dir, args.env, start_time))
else:
state_to_save = player.model.state_dict()
torch.save(state_to_save, '{0}{1}_{2}_best.pt'.format(
args.save_model_dir, args.env, start_time))
reward_sum = 0
player.eps_len = 0
state = player.env.reset()
player.eps_len += 2
#stop after total steps gradient updates have passed
if step >= args.total_frames/args.num_steps:
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
state_to_save = player.model.state_dict()
torch.save(state_to_save, '{0}{1}_{2}_last.pt'.format(
args.save_model_dir, args.env, start_time))
else:
state_to_save = player.model.state_dict()
torch.save(state_to_save, '{0}{1}_{2}_last.pt'.format(
args.save_model_dir, args.env, start_time))
return
time.sleep(10)
player.state = torch.from_numpy(state).float()
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.state = player.state.cuda()
def train(rank, args, shared_models, optimizers, env_conf):
ptitle('Training Agent: {}'.format(rank))
gpu_id = args.gpu_ids[rank % len(args.gpu_ids)]
torch.manual_seed(args.seed + rank)
if gpu_id >= 0:
torch.cuda.manual_seed(args.seed + rank)
env = atari_env(args.env, env_conf, args)
env.seed(args.seed + rank)
player = Agent(env, args, gpu_id)
player.rank = rank
player.state = player.env.reset()
player.state = torch.from_numpy(player.state).float()
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.state = player.state.cuda()
player.models[0].train()
player.models[1].train()
player.eps_len += 2
# player.test_models()
while True:
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
# player.model.load_state_dict(shared_model.state_dict())
player.models[0].load_state_dict(shared_models[0].state_dict())
player.models[1].load_state_dict(shared_models[1].state_dict())
else:
# player.model.load_state_dict(shared_model.state_dict())
player.models[0].load_state_dict(shared_models[0].state_dict())
player.models[1].load_state_dict(shared_models[1].state_dict())
if player.done:
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.cx = Variable(torch.zeros(1, 512).cuda())
player.hx = Variable(torch.zeros(1, 512).cuda())
else:
player.cx = Variable(torch.zeros(1, 512))
player.hx = Variable(torch.zeros(1, 512))
else:
player.cx = Variable(player.cx.data)
player.hx = Variable(player.hx.data)
for step in range(args.num_steps):
player.action_train()
if player.done:
break
# if rank == 0:
# print(player.episodic_reward)
player.episodic_reward = 0
if player.done:
state = player.env.reset()
player.state = torch.from_numpy(state).float()
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.state = player.state.cuda()
R = torch.zeros(1, 1)
if not player.done:
value, _, _ = player.models[player.curr_model_id]((Variable(player.state.unsqueeze(0)),
(player.hx, player.cx)))
R = value.data
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
R = R.cuda()
# player.values.append(Variable(R))
gae = torch.zeros(1, 1)
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
gae = gae.cuda()
R = Variable(R)
# print("Length of values vector", len(player.values))
# print("Length of rewards vector", len(player.rewards))
# print("Length of model sequence vector", len(player.model_sequence))
next_val = Variable(R)
last_val = next_val
R_vec = [Variable(R), Variable(R)]
# last_id = player.model_sequence[-1]
active_flags = [False, False]
policy_loss = [0, 0]
value_loss = [0, 0]
for reward, value, model_id, log_prob, entropy in zip(
reversed(player.rewards),
reversed(player.values),
reversed(player.model_sequence),
reversed(player.log_probs),
reversed(player.entropies)
):
active_flags[model_id] = True
R_vec[model_id] = args.gamma * R_vec[model_id] + reward
R_vec[(model_id+1)%2] *= args.gamma
advantage = R_vec[model_id] - value
value_loss[model_id] += 0.5 * advantage.pow(2)
delta_t = reward + args.gamma * next_val.data - value.data
gae = gae * args.gamma * args.tau + delta_t
policy_loss[model_id] -= (log_prob * Variable(gae) + 0.01 * entropy)
next_val = value
try:
if active_flags[0] is True:
player.models[0].zero_grad()
(policy_loss[0] + 0.5 * value_loss[0]).backward()
ensure_shared_grads(player.models[0], shared_models[0], gpu = gpu_id >= 0)
optimizers[0].step()
if active_flags[1] is True:
player.models[1].zero_grad()
(policy_loss[1] + 0.5 * value_loss[1]).backward()
ensure_shared_grads(player.models[1], shared_models[1], gpu = gpu_id >= 0)
optimizers[1].step()
except Exception as e:
print("Exception caught. Ignoring")
if rank == 1:
print(rewards)
print(model_sequence)
player.clear_actions()
def train(rank, args, shared_model, optimizer, env_conf):
ptitle('Training Agent: {}'.format(rank))
gpu_id = args.gpu_ids[rank % len(args.gpu_ids)]
torch.manual_seed(args.seed + rank)
if gpu_id >= 0:
torch.cuda.manual_seed(args.seed + rank)
env = atari_env(args.env, env_conf, args)
if optimizer is None:
if args.optimizer == 'RMSprop':
optimizer = optim.RMSprop(shared_model.parameters(), lr=args.lr)
if args.optimizer == 'Adam':
optimizer = optim.Adam(shared_model.parameters(),
lr=args.lr,
amsgrad=args.amsgrad)
env.seed(args.seed + rank)
tp_weight = args.tp
player = Agent(None, env, args, None)
player.gpu_id = gpu_id
player.model = A3Clstm(player.env.observation_space.shape[0],
player.env.action_space, args.terminal_prediction,
args.reward_prediction)
player.state = player.env.reset()
player.state = torch.from_numpy(player.state).float()
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.state = player.state.cuda()
player.model = player.model.cuda()
player.model.train()
# Below is where the cores are running episodes continously ...
average_ep_length = 0
while True:
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.model.load_state_dict(shared_model.state_dict())
else:
player.model.load_state_dict(shared_model.state_dict())
if player.done:
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.cx = Variable(torch.zeros(1, 128).cuda())
player.hx = Variable(torch.zeros(1, 128).cuda())
else:
player.cx = Variable(torch.zeros(1, 128))
player.hx = Variable(torch.zeros(1, 128))
else:
player.cx = Variable(player.cx.data)
player.hx = Variable(player.hx.data)
for step in range(args.num_steps):
player.eps_len += 1
player.action_train()
if player.done:
break
if player.done:
state = player.env.reset()
player.state = torch.from_numpy(state).float()
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.state = player.state.cuda()
R = torch.zeros(1, 1)
if not player.done:
value, _, _, _, _ = player.model(
(Variable(player.state.unsqueeze(0)), (player.hx, player.cx)))
R = value.data
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
R = R.cuda()
player.values.append(Variable(R))
policy_loss = 0
value_loss = 0
reward_pred_loss = 0
terminal_loss = 0
gae = torch.zeros(1, 1)
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
gae = gae.cuda()
R = Variable(R) # TODO why this is here?
for i in reversed(range(len(player.rewards))):
R = args.gamma * R + player.rewards[i]
advantage = R - player.values[i]
value_loss = value_loss + 0.5 * advantage.pow(2)
# Generalized Advantage Estimataion
delta_t = player.rewards[i] + args.gamma * player.values[
i + 1].data - player.values[i].data
gae = gae * args.gamma * args.tau + delta_t
policy_loss = policy_loss - player.log_probs[i] * Variable(
gae) - 0.01 * player.entropies[i]
if args.reward_prediction:
reward_pred_loss = reward_pred_loss + (
player.reward_predictions[i] - player.rewards[i]).pow(2)
if args.terminal_prediction: # new way of using emprical episode length as a proxy for current length.
if player.average_episode_length is None:
end_predict_labels = np.arange(
player.eps_len - len(player.terminal_predictions),
player.eps_len) / player.eps_len # heuristic
else:
end_predict_labels = np.arange(
player.eps_len - len(player.terminal_predictions),
player.eps_len) / player.average_episode_length
for i in range(len(player.terminal_predictions)):
terminal_loss = terminal_loss + (
player.terminal_predictions[i] -
end_predict_labels[i]).pow(2)
terminal_loss = terminal_loss / len(player.terminal_predictions)
player.model.zero_grad()
#print(f"policy loss {policy_loss} and value loss {value_loss} and terminal loss {terminal_loss} and reward pred loss {reward_pred_loss}")
total_loss = policy_loss + 0.5 * value_loss + tp_weight * terminal_loss + 0.5 * reward_pred_loss
total_loss.backward() # will free memory ...
# Visualize Computation Graph
#graph = make_dot(total_loss)
#from graphviz import Source
#Source.view(graph)
ensure_shared_grads(player.model, shared_model, gpu=gpu_id >= 0)
optimizer.step()
player.clear_actions()
if player.done:
if player.average_episode_length is None: # initial one
player.average_episode_length = player.eps_len
else:
player.average_episode_length = int(
0.99 * player.average_episode_length +
0.01 * player.eps_len)
#print(player.average_episode_length, 'current one is ', player.eps_len)
player.eps_len = 0 # reset here
def train_func(rank, args, shared_model, optimizer, env_conf, datasets):
if args.deploy:
return
ptitle('Train {0}'.format(rank))
print('Start training agent: ', rank)
if rank == 0:
logger = Logger(args.log_dir[:-1] + '_losses/')
train_step = 0
gpu_id = args.gpu_ids[rank % len(args.gpu_ids)]
env_conf["env_gpu"] = gpu_id
torch.manual_seed(args.seed + rank)
if gpu_id >= 0:
torch.cuda.manual_seed(args.seed + rank)
env = Debug_env(datasets, env_conf, seed=args.seed + rank)
if optimizer is None:
if args.optimizer == 'RMSprop':
optimizer = optim.RMSprop(shared_model.parameters(), lr=args.lr)
if args.optimizer == 'Adam':
optimizer = optim.Adam(shared_model.parameters(),
lr=args.lr,
amsgrad=args.amsgrad)
player = Agent(None, env, args, None)
player.gpu_id = gpu_id
nChan = 3
if args.is3D:
nChan = 4
if args.alpha_only:
nChan = 1
if not args.is3D:
player.model = get_model(args,
"ENet",
input_shape=env_conf["obs_shape"],
num_actions=args.num_actions * nChan)
elif not args.obs3D:
player.model = get_model(args,
"ENet",
input_shape=env_conf["obs_shape"],
num_actions=args.num_actions * nChan)
elif args.obs3D:
player.model = get_model(args,
"Net3D",
input_shape=env_conf["obs_shape"],
num_actions=args.num_actions * nChan)
player.state = player.env.reset()
player.state = torch.from_numpy(player.state).float()
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.state = player.state.cuda()
player.model = player.model.cuda()
player.model.train()
if rank == 0:
eps_reward = 0
pinned_eps_reward = 0
while True:
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.model.load_state_dict(shared_model.state_dict())
else:
player.model.load_state_dict(shared_model.state_dict())
if player.done:
player.eps_len = 0
if rank == 0:
if train_step % args.train_log_period == 0 and train_step > 0:
print("train: step", train_step, "\teps_reward",
eps_reward)
if train_step > 0:
pinned_eps_reward = player.env.sum_rewards.mean()
eps_reward = 0
for step in range(args.num_steps):
player.action_train()
if rank == 0:
eps_reward = player.env.sum_rewards.mean()
if player.done:
break
if player.done:
state = player.env.reset()
player.state = torch.from_numpy(state).float()
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.state = player.state.cuda()
if not args.alpha_only:
if not args.is3D:
R = torch.zeros(1, 1, args.num_actions * 3)
else:
R = torch.zeros(1, 1, args.num_actions * 4)
else:
R = torch.zeros(1, 1, args.num_actions)
if not player.done:
value, _ = player.model(Variable(player.state.unsqueeze(0)))
R = value.data
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
R = R.cuda()
player.values.append(Variable(R))
policy_loss = 0
value_loss = 0
if not args.alpha_only:
if not args.is3D:
gae = torch.zeros(1, 1, args.num_actions * 3)
else:
gae = torch.zeros(1, 1, args.num_actions * 4)
else:
gae = torch.zeros(1, 1, args.num_actions)
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
gae = gae.cuda()
R = Variable(R)
for i in reversed(range(len(player.rewards))):
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
reward_i = torch.tensor(player.rewards[i]).cuda()
else:
reward_i = torch.tensor(player.rewards[i])
R = args.gamma * R + reward_i
advantage = R - player.values[i]
value_loss = value_loss + (0.5 * advantage * advantage).mean()
delta_t = player.values[
i + 1].data * args.gamma + reward_i - player.values[i].data
gae = gae * args.gamma * args.tau + delta_t
policy_loss = policy_loss - \
(player.log_probs[i] * Variable(gae)).mean () - \
(args.entropy_alpha * player.entropies[i]).mean ()
player.model.zero_grad()
sum_loss = (policy_loss + value_loss)
curtime = time.time()
sum_loss.backward()
ensure_shared_grads(player.model, shared_model, gpu=gpu_id >= 0)
curtime = time.time()
optimizer.step()
player.clear_actions()
if rank == 0:
train_step += 1
if train_step % args.log_period * 10 == 0 and train_step > 0:
log_info = {
'train: value_loss': value_loss,
'train: policy_loss': policy_loss,
'train: eps reward': pinned_eps_reward,
}
for tag, value in log_info.items():
logger.scalar_summary(tag, value, train_step)
def train_robust(rank, args, shared_model, optimizer, env_conf):
ptitle('Training Agent: {}'.format(rank))
gpu_id = args.gpu_ids[rank % len(args.gpu_ids)]
if args.seed:
torch.manual_seed(args.seed + rank)
if gpu_id >= 0:
torch.cuda.manual_seed(args.seed + rank)
env = atari_env(args.env, env_conf, args)
if optimizer is None:
if args.optimizer == 'RMSprop':
optimizer = optim.RMSprop(shared_model.parameters(), lr=args.lr)
if args.optimizer == 'Adam':
optimizer = optim.Adam(shared_model.parameters(),
lr=args.lr,
amsgrad=args.amsgrad)
if args.seed:
env.seed(args.seed + rank)
player = Agent(None, env, args, None)
player.gpu_id = gpu_id
player.model = A3Cff(player.env.observation_space.shape[0],
player.env.action_space)
player.state = player.env.reset()
player.state = torch.from_numpy(player.state).float()
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.state = player.state.cuda()
player.model = player.model.cuda()
player.model.train()
player.eps_len += 2
while True:
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.model.load_state_dict(shared_model.state_dict())
else:
player.model.load_state_dict(shared_model.state_dict())
p = optimizer.param_groups[0]['params'][0]
step = optimizer.state[p]['step']
if step >= (args.total_frames / args.num_steps):
return
#increase linearly until 2/3 through halfway
lin_coeff = min(1, (1.5 * int(step) + 1) /
(args.total_frames / args.num_steps))
epsilon = lin_coeff * args.epsilon_end
kappa = args.kappa_end #(1-lin_coeff)*1 + lin_coeff*args.kappa_end
for step in range(args.num_steps):
player.action_train(bound_epsilon=epsilon)
if player.done:
break
if player.done:
state = player.env.reset()
player.state = torch.from_numpy(state).float()
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.state = player.state.cuda()
R = torch.zeros(1, 1)
if not player.done:
value, _ = player.model(Variable(player.state.unsqueeze(0)))
R = value.data
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
R = R.cuda()
player.values.append(Variable(R))
policy_loss = 0
value_loss = 0
gae = torch.zeros(1, 1)
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
gae = gae.cuda()
R = Variable(R)
for i in reversed(range(len(player.rewards))):
R = args.gamma * R + player.rewards[i]
advantage = R - player.values[i]
value_loss = value_loss + 0.5 * advantage.pow(2)
# Generalized Advantage Estimataion
delta_t = player.rewards[i] + args.gamma * \
player.values[i + 1].data - player.values[i].data
gae = gae * args.gamma * args.tau + delta_t
if gae >= 0:
worst_case_loss = -player.min_log_probs[i] * Variable(gae)
else:
worst_case_loss = -player.max_log_probs[i] * Variable(gae)
standard_loss = -player.log_probs[i] * Variable(gae)
policy_loss = policy_loss + kappa * standard_loss + (
1 - kappa) * worst_case_loss - 0.01 * player.entropies[i]
#print(policy_loss + 0.5 * value_loss)
player.model.zero_grad()
(policy_loss + 0.5 * value_loss).backward()
ensure_shared_grads(player.model, shared_model, gpu=gpu_id >= 0)
optimizer.step()
player.clear_actions()
def train(rank, args, shared_model, optimizer, env_conf):
ptitle('Training Agent: {}'.format(rank))
gpu_id = args.gpu_ids[rank % len(args.gpu_ids)]
torch.manual_seed(args.seed + rank)
if gpu_id >= 0:
torch.cuda.manual_seed(args.seed + rank)
env = Environment() # 創建環境
if optimizer is None:
if args.optimizer == 'RMSprop':
optimizer = optim.RMSprop(shared_model.parameters(), lr=args.lr)
if args.optimizer == 'Adam':
optimizer = optim.Adam(shared_model.parameters(),
lr=args.lr,
amsgrad=args.amsgrad)
# env.seed(args.seed + rank)
player = Agent(None, env, args, None) # 創建代理人
player.gpu_id = gpu_id
num_actions = env.get_num_actions()
player.model = A3Clstm(
Config.STACKED_FRAMES, # A3C模型
num_actions)
player.state, available = player.env.reset() # 初始環境
player.state = torch.from_numpy(player.state).float()
player.available = torch.from_numpy(available).float()
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.state = player.state.cuda()
player.model = player.model.cuda()
player.available = player.available.cuda()
player.model.train() # 訓練模式
player.eps_len += 1
while True:
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.model.load_state_dict(shared_model.state_dict())
else:
player.model.load_state_dict(shared_model.state_dict()) # 更新網路
if player.done:
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.cx = Variable(torch.zeros(1, 512).cuda())
player.hx = Variable(torch.zeros(1, 512).cuda())
else:
player.cx = Variable(torch.zeros(1, 512))
player.hx = Variable(torch.zeros(1, 512)) # 完成一次訓練 初始化LSTM
else:
player.cx = Variable(player.cx.data)
player.hx = Variable(player.hx.data)
for step in range(args.num_steps): # T-max = 20
player.action_train()
if player.done:
break
if player.done:
state, available = player.env.reset()
player.state = torch.from_numpy(state).float()
player.available = torch.from_numpy(available).float()
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.state = player.state.cuda()
player.available = player.available.cuda()
R = torch.zeros(1, 1) # if done : R_t-max = 0
if not player.done:
value, _, _, _ = player.model(
(Variable(player.state.unsqueeze(0)), (player.hx, player.cx)))
R = value.data # R_t-max = V(s)
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
R = R.cuda()
player.values.append(Variable(R))
policy_loss = 0
value_loss = 0
gae = torch.zeros(1, 1)
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
gae = gae.cuda()
R = Variable(R)
for i in reversed(range(len(player.rewards))):
R = args.gamma * R + player.rewards[i]
advantage = R - player.values[i]
value_loss = value_loss + 0.5 * advantage.pow(2)
# Generalized Advantage Estimataion
delta_t = player.rewards[i] + args.gamma * \
player.values[i + 1].data - player.values[i].data
gae = gae * args.gamma * args.tau + delta_t
policy_loss = policy_loss - \
player.log_probs[i] * \
Variable(gae) - 0.01 * player.entropies[i]
player.model.zero_grad()
(policy_loss + 0.5 * value_loss).backward()
ensure_shared_grads(player.model, shared_model, gpu=gpu_id >= 0)
optimizer.step()
player.clear_actions()
def train(rank, args, shared_model, optimizer, env_conf, iters,
checkpoint_path):
iters = dill.loads(iters)
if args.enable_gavel_iterator and rank == 0:
iters._init_logger()
ptitle('Training Agent: {}'.format(rank))
gpu_id = args.gpu_ids[rank % len(args.gpu_ids)]
torch.manual_seed(args.seed + rank)
if gpu_id >= 0:
torch.cuda.manual_seed(args.seed + rank)
env = atari_env(args.env, env_conf, args)
if optimizer is None:
if args.optimizer == 'RMSprop':
optimizer = optim.RMSprop(shared_model.parameters(), lr=args.lr)
if args.optimizer == 'Adam':
optimizer = optim.Adam(shared_model.parameters(),
lr=args.lr,
amsgrad=args.amsgrad)
env.seed(args.seed + rank)
player = Agent(None, env, args, None)
player.gpu_id = gpu_id
player.model = A3Clstm(player.env.observation_space.shape[0],
player.env.action_space)
player.state = player.env.reset()
player.state = torch.from_numpy(player.state).float()
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.state = player.state.cuda()
player.model = player.model.cuda()
player.model.train()
player.eps_len += 2
elapsed_time = 0
start_time = time.time()
for i in iters:
if i % 100 == 0:
print('GPU %d finished step %d' % (rank, i), flush=True)
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.model.load_state_dict(shared_model.state_dict())
else:
player.model.load_state_dict(shared_model.state_dict())
if player.done:
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.cx = Variable(torch.zeros(1, 512).cuda())
player.hx = Variable(torch.zeros(1, 512).cuda())
else:
player.cx = Variable(torch.zeros(1, 512))
player.hx = Variable(torch.zeros(1, 512))
else:
player.cx = Variable(player.cx.data)
player.hx = Variable(player.hx.data)
for step in range(args.num_steps):
player.action_train()
if player.done:
break
if player.done:
state = player.env.reset()
player.state = torch.from_numpy(state).float()
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.state = player.state.cuda()
R = torch.zeros(1, 1)
if not player.done:
value, _, _ = player.model(
(Variable(player.state.unsqueeze(0)), (player.hx, player.cx)))
R = value.data
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
R = R.cuda()
player.values.append(Variable(R))
policy_loss = 0
value_loss = 0
gae = torch.zeros(1, 1)
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
gae = gae.cuda()
R = Variable(R)
for i in reversed(range(len(player.rewards))):
R = args.gamma * R + player.rewards[i]
advantage = R - player.values[i]
value_loss = value_loss + 0.5 * advantage.pow(2)
# Generalized Advantage Estimataion
delta_t = player.rewards[i] + args.gamma * \
player.values[i + 1].data - player.values[i].data
gae = gae * args.gamma * args.tau + delta_t
policy_loss = policy_loss - \
player.log_probs[i] * \
Variable(gae) - 0.01 * player.entropies[i]
player.model.zero_grad()
(policy_loss + 0.5 * value_loss).backward()
ensure_shared_grads(player.model, shared_model, gpu=gpu_id >= 0)
optimizer.step()
player.clear_actions()
elapsed_time += time.time() - start_time
start_time = time.time()
if (args.throughput_estimation_interval is not None
and i % args.throughput_estimation_interval == 0
and rank == 0):
print('[THROUGHPUT_ESTIMATION]\t%s\t%d' % (time.time(), i))
if (args.max_duration is not None
and elapsed_time >= args.max_duration):
break
if args.enable_gavel_iterator and rank == 0:
state = shared_model.state_dict()
iters.save_checkpoint(state, checkpoint_path)
iters.complete()
def train(rank, args, shared_model, optimizer, train_modes, n_iters, env=None):
n_iter = 0
writer = SummaryWriter(os.path.join(args.log_dir, 'Agent:{}'.format(rank)))
ptitle('Training Agent: {}'.format(rank))
gpu_id = args.gpu_ids[rank % len(args.gpu_ids)]
torch.manual_seed(args.seed + rank)
training_mode = args.train_mode
env_name = args.env
train_modes.append(training_mode)
n_iters.append(n_iter)
if gpu_id >= 0:
torch.cuda.manual_seed(args.seed + rank)
device = torch.device('cuda:' + str(gpu_id))
if len(args.gpu_ids) > 1:
device_share = torch.device('cpu')
else:
device_share = torch.device('cuda:' + str(args.gpu_ids[-1]))
else:
device = device_share = torch.device('cpu')
if env is None:
env = create_env(env_name, args)
if args.train_mode == 0:
params = shared_model.player0.parameters()
elif args.train_mode == 1:
params = shared_model.player1.parameters()
else:
params = shared_model.parameters()
if optimizer is None:
if args.optimizer == 'RMSprop':
optimizer = optim.RMSprop(params, lr=args.lr)
if args.optimizer == 'Adam':
optimizer = optim.Adam(params, lr=args.lr)
env.seed(args.seed)
player = Agent(None, env, args, None, device)
player.w_entropy_target = args.entropy_target
player.gpu_id = gpu_id
# prepare model
player.model = build_model(player.env.observation_space,
player.env.action_space, args, device)
player.model = player.model.to(device)
player.model.train()
player.reset()
reward_sum = torch.zeros(player.num_agents).to(device)
reward_sum_org = np.zeros(player.num_agents)
ave_reward = np.zeros(2)
ave_reward_longterm = np.zeros(2)
count_eps = 0
while True:
# sys to the shared model
player.model.load_state_dict(shared_model.state_dict())
if player.done:
player.reset()
reward_sum = torch.zeros(player.num_agents).to(device)
reward_sum_org = np.zeros(player.num_agents)
count_eps += 1
player.update_rnn_hiden()
t0 = time.time()
for i in range(args.num_steps):
player.action_train()
reward_sum += player.reward
reward_sum_org += player.reward_org
if player.done:
for j, r_i in enumerate(reward_sum):
writer.add_scalar('train/reward_' + str(j), r_i,
player.n_steps)
break
fps = i / (time.time() - t0)
# cfg training mode
# 0: tracker 1: target -1:joint all
training_mode = train_modes[rank]
policy_loss, value_loss, entropies, pred_loss = player.optimize(
params, optimizer, shared_model, training_mode, device_share)
for i in range(min(player.num_agents, 3)):
writer.add_scalar('train/policy_loss_' + str(i),
policy_loss[i].mean(), player.n_steps)
writer.add_scalar('train/value_loss_' + str(i), value_loss[i],
player.n_steps)
writer.add_scalar('train/entropies' + str(i), entropies[i].mean(),
player.n_steps)
writer.add_scalar('train/pred_R_loss', pred_loss, player.n_steps)
writer.add_scalar('train/ave_reward',
ave_reward[0] - ave_reward_longterm[0],
player.n_steps)
writer.add_scalar('train/mode', training_mode, player.n_steps)
writer.add_scalar('train/fps', fps, player.n_steps)
n_iter += 1
n_iters[rank] = n_iter
if train_modes[rank] == -100:
env.close()
break
def train(rank, args, shared_model, optimizer, env_conf, shared_counter,
targ_shared):
ptitle('Training Agent: {}'.format(rank))
gpu_id = args.gpu_ids[rank % len(args.gpu_ids)]
device = torch.device('cuda:{}'.format(gpu_id) if gpu_id >= 0 else 'cpu')
torch.manual_seed(args.seed + rank)
torch.cuda.manual_seed(args.seed + rank)
env = atari_env(args.env, env_conf, args)
if optimizer is None:
if args.optimizer == 'RMSprop':
optimizer = optim.RMSprop(shared_model.parameters(), lr=args.lr)
if args.optimizer == 'Adam':
optimizer = optim.Adam(shared_model.parameters(),
lr=args.lr,
amsgrad=args.amsgrad)
env.seed(args.seed + rank)
player = Agent(None, env, args, None, gpu_id=gpu_id)
player.model = A3Clstm(player.env.observation_space.shape[0],
player.env.action_space)
player.model.apply(weights_init)
player.state = player.env.reset()
player.state = torch.from_numpy(player.state).to(torch.float32)
player.state = player.state.to(device)
player.model = player.model.to(device)
#player.targ_model = copy.deepcopy(player.model)
player.model.train()
#player.targ_model.eval()
player.eps_len += 2
while True:
player.model.load_state_dict(shared_model.state_dict())
#player.targ_model.load_state_dict(targ_shared.state_dict())
if player.done:
player.cx = torch.zeros(1, 512).to(device)
player.hx = torch.zeros(1, 512).to(device)
#player.targ_cx = copy.deepcopy(player.cx).detach()
#player.targ_hx = copy.deepcopy(player.hx).detach()
else:
player.cx = player.cx.detach()
player.hx = player.hx.detach()
for step in range(args.num_steps):
player.action_train()
if player.done:
break
if player.done:
state = player.env.reset()
player.state = torch.from_numpy(state).to(torch.float32)
player.state = player.state.to(device)
#alpha = player.model.log_alpha.exp().detach()
alpha = .01
#alpha = 0
x_R = torch.zeros(1, 1)
if not player.done:
with torch.no_grad():
action, value, logit, q_value, _ = player.model(
(player.state.unsqueeze(0), (player.hx, player.cx)))
x_R = q_value[1].detach() - alpha * F.log_softmax(
logit, -1).gather(-1, action)
x_R = x_R.to(device)
policy_loss = 0
adv_gae_loss = 0
for i in reversed(range(len(player.rewards))):
x_R = args.gamma * x_R + player.rewards[i]
adv_gae_loss = adv_gae_loss + (player.tra_adv_gae[i][1] -
x_R.detach()).pow(2) * .5
#policy_loss = policy_loss - player.log_probs[i] * player.tra_adv_gae[i][0].detach() + alpha * player.log_probs[i] * player.log_probs[i].detach()
policy_loss = policy_loss - (F.softmax(
player.values[i], -1) * player.tra_adv_gae[i][0].detach()).sum(
-1) - alpha * player.entropies[i].unsqueeze(-1)
#policy_loss = policy_loss - player.log_probs[i] * (x_R - (F.softmax(player.values[i], -1) *
# player.tra_adv_gae[i][0]).sum(-1) - alpha * player.entropies[i]).detach() + alpha * player.log_probs[i] * player.log_probs[i].detach()
#prob = F.softmax(player.values[i], -1)
#ent_alpha = alpha * player.entropies[i].unsqueeze(-1)
#advs = (player.tra_adv_gae[i][0] -
# ((player.tra_adv_gae[i][0] * prob).sum(-1, True) +
# ent_alpha)).detach()
#policy_loss = policy_loss - (prob * advs).sum(-1) - ent_alpha
x_R = x_R - alpha * player.log_probs[i].detach()
player.model.zero_grad()
(policy_loss + .5 * adv_gae_loss).backward(retain_graph=False)
ensure_shared_grads(player.model, shared_model, gpu=gpu_id >= 0)
optimizer.step()
player.clear_actions()
with shared_counter.get_lock():
shared_counter.value += len(player.rewards)
if shared_counter.value > args.interact_steps:
break
def train(rank, args, shared_model, optimizer, optimizer_r, env_conf, lock,
counter):
ptitle('Training Agent: {}'.format(rank))
gpu_id = args.gpu_ids[rank % len(args.gpu_ids)]
torch.manual_seed(args.seed + rank)
if gpu_id >= 0:
torch.cuda.manual_seed(args.seed + rank)
env = atari_env(args.env, env_conf, args)
if optimizer is None:
if args.optimizer == 'RMSprop':
optimizer = optim.RMSprop(shared_model.parameters(), lr=args.lr)
if args.optimizer == 'Adam':
optimizer = optim.Adam(shared_model.parameters(),
lr=args.lr,
amsgrad=args.amsgrad)
env.seed(args.seed + rank)
player = Agent(None, env, args, None)
player.gpu_id = gpu_id
player.model = A3Clstm(player.env.observation_space.shape[0],
player.env.action_space)
player.state = player.env.reset()
player.state = torch.from_numpy(player.state).float()
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.state = player.state.cuda()
player.model = player.model.cuda()
player.model.train()
player.eps_len += 2
while True:
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.model.load_state_dict(shared_model.state_dict())
else:
player.model.load_state_dict(shared_model.state_dict())
if player.done:
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.cx = [
Variable(torch.zeros(1, 512).cuda()),
Variable(torch.zeros(1, 512).cuda())
]
player.hx = [
Variable(torch.zeros(1, 512).cuda()),
Variable(torch.zeros(1, 512).cuda())
]
else:
player.cx = [
Variable(torch.zeros(1, 512)),
Variable(torch.zeros(1, 512))
]
player.hx = [
Variable(torch.zeros(1, 512)),
Variable(torch.zeros(1, 512))
]
else:
player.cx = [
Variable(player.cx[0].data),
Variable(player.cx[1].data)
]
player.hx = [
Variable(player.hx[0].data),
Variable(player.cx[1].data)
]
# 测试rnet的更新有没有影响到这里
# ps = list(player.model.r_net.named_parameters())
# n, v = ps[6]
# print(v.sum())
for step in range(args.num_steps):
player.action_train()
if player.done:
break
if player.done:
state = player.env.reset()
player.state = torch.from_numpy(state).float()
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.state = player.state.cuda()
R = torch.zeros(1, 1)
if not player.done:
value, _, _, _ = player.model(
(Variable(player.state.unsqueeze(0)),
(player.hx[0], player.cx[0]), (player.hx[1], player.cx[1])))
R = value.data
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
R = R.cuda()
player.values.append(Variable(R))
policy_loss = 0
value_loss = 0
gae = torch.zeros(1, 1)
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
gae = gae.cuda()
R = Variable(R)
for i in reversed(range(len(player.rewards))):
R = args.gamma * R + player.rewards[i]
advantage = R - player.values[i]
value_loss = value_loss + 0.5 * advantage.pow(2)
# Generalized Advantage Estimataion
delta_t = player.rewards[i] + args.gamma * \
player.values[i + 1].data - player.values[i].data
gae = gae * args.gamma * args.tau + delta_t
policy_loss = policy_loss - \
player.log_probs[i] * \
Variable(gae) - 0.01 * player.entropies[i]
with lock:
counter.value += 1
# rnet
player.model.r_net.zero_grad()
(args.actor_weight * policy_loss +
(1 - args.actor_weight) * value_loss).backward(retain_graph=True)
ensure_shared_grads(player.model.r_net,
shared_model.r_net,
gpu=gpu_id >= 0)
optimizer_r.step()
player.model.zero_grad()
(policy_loss + 0.5 * value_loss).backward()
player.model.r_net.zero_grad()
ensure_shared_grads(player.model, shared_model, gpu=gpu_id >= 0)
optimizer.step()
player.clear_actions()
def train(rank, args, shared_model, optimizer, env_conf, num_tau_samples=32, num_tau_prime_samples=32, kappa=1.0, num_quantiles=32):
ptitle('Training Agent: {}'.format(rank))
gpu_id = args.gpu_ids[rank % len(args.gpu_ids)]
torch.manual_seed(args.seed + rank)
if gpu_id >= 0:
torch.cuda.manual_seed(args.seed + rank)
env = atari_env(args.env, env_conf, args)
if optimizer is None:
if args.optimizer == 'RMSprop':
optimizer = optim.RMSprop(shared_model.parameters(), lr=args.lr)
if args.optimizer == 'Adam':
optimizer = optim.Adam(
shared_model.parameters(), lr=args.lr, amsgrad=args.amsgrad)
env.seed(args.seed + rank)
player = Agent(None, env, args, None)
player.gpu_id = gpu_id
player.model = A3Clstm(player.env.observation_space.shape[0],
player.env.action_space)
player.state = player.env.reset()
player.state = torch.from_numpy(player.state).float()
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.state = player.state.cuda()
player.model = player.model.cuda()
player.model.train()
player.eps_len += 2
while True:
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.model.load_state_dict(shared_model.state_dict())
else:
player.model.load_state_dict(shared_model.state_dict())
if player.done:
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.cx = Variable(torch.zeros(1, 512).cuda())
player.hx = Variable(torch.zeros(1, 512).cuda())
else:
player.cx = Variable(torch.zeros(1, 512))
player.hx = Variable(torch.zeros(1, 512))
else:
player.cx = Variable(player.cx.data)
player.hx = Variable(player.hx.data)
for step in range(args.num_steps):
player.action_train()
if player.done:
break
if player.done:
state = player.env.reset()
player.state = torch.from_numpy(state).float()
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.state = player.state.cuda()
R = torch.zeros(1,num_tau_prime_samples)
if not player.done:
logit, _, _ = player.model((Variable(
player.state.unsqueeze(0)), (player.hx, player.cx)))
q_vals = torch.mean(logit,0)
_, action = torch.max(q_vals,0)
logit, _, _ = player.model((Variable(player.state.unsqueeze(0)),
(player.hx, player.cx)))
R = logit[:,action]
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
R = R.cuda()
#R = R.detach()
R = Variable(R)
value_loss = 0
for i in reversed(range(len(player.rewards))):
R = args.gamma * R + player.rewards[i]
advantage = R.repeat(num_tau_samples,1) - player.logits_array[i].repeat(1, num_tau_prime_samples)
#print("Ad: ",advantage)
loss = (torch.abs(advantage) <= kappa).float() * 0.5 * advantage ** 2
#print("loss: ",loss.sum(0).sum(0), loss)
loss += (torch.abs(advantage) > kappa).float() * kappa * (torch.abs(advantage) - 0.5 * kappa)
#print("loss: ",loss.sum(0).sum(0), loss)
step_loss = torch.abs(player.quantiles_array[i].cuda() - (advantage.detach()<0).float()) * loss/kappa
value_loss += step_loss.sum(0).mean(0)
player.model.zero_grad()
value_loss.backward()
ensure_shared_grads(player.model, shared_model, gpu=gpu_id >= 0)
optimizer.step()
player.clear_actions()
def train_func(rank,
args,
shared_model,
optimizer,
env_conf,
datasets=None,
shared_dict=None):
if args.deploy:
return
ptitle('Train {0}'.format(rank))
print('Start training agent: ', rank)
if rank == 0:
logger = Logger(args.log_dir[:-1] + '_losses/')
train_step = 0
gpu_id = args.gpu_ids[rank % len(args.gpu_ids)]
env_conf["env_gpu"] = gpu_id
torch.manual_seed(args.seed + rank)
if gpu_id >= 0:
torch.cuda.manual_seed(args.seed + rank)
raw_list, gt_lbl_list = datasets
env = EM_env(raw_list,
env_conf,
type="train",
gt_lbl_list=gt_lbl_list,
seed=args.seed + rank)
if optimizer is None:
if args.optimizer == 'RMSprop':
optimizer = optim.RMSprop(shared_model.parameters(), lr=args.lr)
if args.optimizer == 'Adam':
optimizer = optim.Adam(shared_model.parameters(),
lr=args.lr,
amsgrad=args.amsgrad)
player = Agent(None, env, args, None)
player.gpu_id = gpu_id
player.model = get_model(args,
args.model,
env.observation_space.shape,
args.features,
atrous_rates=args.atr_rate,
num_actions=2,
split=args.data_channel,
gpu_id=gpu_id,
multi=args.multi)
player.state = player.env.reset()
player.state = torch.from_numpy(player.state).float()
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.state = player.state.cuda()
player.model = player.model.cuda()
player.model.train()
if rank == 0:
eps_reward = 0
pinned_eps_reward = 0
while True:
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.model.load_state_dict(shared_model.state_dict())
else:
player.model.load_state_dict(shared_model.state_dict())
if player.done:
player.eps_len = 0
if rank == 0:
if train_step % args.train_log_period == 0 and train_step > 0:
print("train: step", train_step, "\teps_reward",
eps_reward)
if train_step > 0:
pinned_eps_reward = player.env.sum_reward.mean()
eps_reward = 0
if args.lstm_feats:
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.cx, player.hx = player.model.lstm.init_hidden(
batch_size=1, use_cuda=True)
else:
player.cx, player.hx = player.model.lstm.init_hidden(
batch_size=1, use_cuda=False)
elif args.lstm_feats:
player.cx = Variable(player.cx.data)
player.hx = Variable(player.hx.data)
for step in range(args.num_steps):
if rank < args.lbl_agents:
player.action_train(use_lbl=True)
else:
player.action_train()
if rank == 0:
eps_reward = player.env.sum_reward.mean()
if player.done:
break
if player.done:
state = player.env.reset(player.model, gpu_id)
player.state = torch.from_numpy(state).float()
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.state = player.state.cuda()
if "3D" in args.data:
R = torch.zeros(1, 1, env_conf["size"][0], env_conf["size"][1],
env_conf["size"][2])
else:
R = torch.zeros(1, 1, env_conf["size"][0], env_conf["size"][1])
if args.lowres:
R = torch.zeros(1, 1, env_conf["size"][0] // 2,
env_conf["size"][1] // 2)
if not player.done:
if args.lstm_feats:
value, _, _ = player.model(
(Variable(player.state.unsqueeze(0)), (player.hx,
player.cx)))
else:
value, _ = player.model(Variable(player.state.unsqueeze(0)))
R = value.data
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
R = R.cuda()
player.values.append(Variable(R))
policy_loss = 0
value_loss = 0
if "3D" in args.data:
gae = torch.zeros(1, 1, env_conf["size"][0], env_conf["size"][1],
env_conf["size"][2])
else:
gae = torch.zeros(1, 1, env_conf["size"][0], env_conf["size"][1])
if args.rew_drop:
keep_map = torch.tensor(player.env.keep_map)
if args.lowres:
gae = torch.zeros(1, 1, env_conf["size"][0] // 2,
env_conf["size"][1] // 2)
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
gae = gae.cuda()
if args.rew_drop:
keep_map = keep_map.cuda()
R = Variable(R)
for i in reversed(range(len(player.rewards))):
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
reward_i = torch.tensor(player.rewards[i]).cuda()
else:
reward_i = torch.tensor(player.rewards[i])
R = args.gamma * R + reward_i
if args.rew_drop:
advantage = R - player.values[i]
value_loss = value_loss + (0.5 * advantage * advantage *
keep_map).mean()
delta_t = player.values[
i + 1].data * args.gamma + reward_i - player.values[i].data
gae = gae * args.gamma * args.tau + delta_t
else:
advantage = R - player.values[i]
value_loss = value_loss + (0.5 * advantage * advantage).mean()
delta_t = player.values[
i + 1].data * args.gamma + reward_i - player.values[i].data
gae = gae * args.gamma * args.tau + delta_t
if args.noisy:
policy_loss = policy_loss - \
(player.log_probs[i] * Variable(gae)).mean ()
else:
if args.rew_drop:
policy_loss = policy_loss - \
(player.log_probs[i] * Variable(gae) * keep_map).mean () - \
(args.entropy_alpha * player.entropies[i] * keep_map).mean ()
else:
policy_loss = policy_loss - \
(player.log_probs[i] * Variable(gae)).mean () - \
(args.entropy_alpha * player.entropies[i]).mean ()
player.model.zero_grad()
sum_loss = (policy_loss + value_loss)
curtime = time.time()
# print ("backward curtime:", curtime)
sum_loss.backward()
# print ("backward done", time.time () - curtime)
ensure_shared_grads(player.model, shared_model, gpu=gpu_id >= 0)
curtime = time.time()
# print ("optim curtime:", curtime)
optimizer.step()
# print ("optim done", time.time () - curtime)
player.clear_actions()
if args.wctrl == "s2m":
player.env.config["spl_w"] = shared_dict["spl_w"]
player.env.config["mer_w"] = shared_dict["mer_w"]
if rank == 0:
train_step += 1
if train_step % args.log_period == 0 and train_step > 0:
log_info = {
'train: value_loss': value_loss,
'train: policy_loss': policy_loss,
'train: eps reward': pinned_eps_reward,
}
if "EX" in args.model:
log_info["cell_prob_loss"] = cell_prob_loss
for tag, value in log_info.items():
logger.scalar_summary(tag, value, train_step)
def train(rank, args, shared_model, optimizer, env_conf):
torch.manual_seed(args.seed + rank)
env = atari_env(args.env, env_conf)
if optimizer is None:
if args.optimizer == 'RMSprop':
optimizer = optim.RMSprop(shared_model.parameters(), lr=args.lr)
if args.optimizer == 'Adam':
optimizer = optim.Adam(shared_model.parameters(), lr=args.lr)
env.seed(args.seed + rank)
player = Agent(None, env, args, None)
player.model = A3Clstm(
player.env.observation_space.shape[0], player.env.action_space)
player.state = player.env.reset()
player.state = torch.from_numpy(player.state).float()
player.model.train()
while True:
player.model.load_state_dict(shared_model.state_dict())
for step in range(args.num_steps):
player.action_train()
if args.count_lives:
player.check_state()
if player.done:
break
if player.done:
player.eps_len = 0
player.current_life = 0
state = player.env.reset()
player.state = torch.from_numpy(state).float()
R = torch.zeros(1, 1)
if not player.done:
value, _, _ = player.model(
(Variable(player.state.unsqueeze(0)), (player.hx, player.cx)))
R = value.data
player.values.append(Variable(R))
policy_loss = 0
value_loss = 0
R = Variable(R)
gae = torch.zeros(1, 1)
for i in reversed(range(len(player.rewards))):
R = args.gamma * R + player.rewards[i]
advantage = R - player.values[i]
value_loss = value_loss + 0.5 * advantage.pow(2)
# Generalized Advantage Estimataion
delta_t = player.rewards[i] + args.gamma * \
player.values[i + 1].data - player.values[i].data
gae = gae * args.gamma * args.tau + delta_t
policy_loss = policy_loss - \
player.log_probs[i] * \
Variable(gae) - 0.01 * player.entropies[i]
optimizer.zero_grad()
(policy_loss + 0.5 * value_loss).backward()
torch.nn.utils.clip_grad_norm(player.model.parameters(), 40)
ensure_shared_grads(player.model, shared_model)
optimizer.step()
player.clear_actions()
def train (rank, args, shared_model, optimizer, env_conf, datasets=None):
ptitle('Training Agent: {}'.format(rank))
print ('Start training agent: ', rank)
if rank == 0:
logger = Logger (args.log_dir)
train_step = 0
gpu_id = args.gpu_ids[rank % len(args.gpu_ids)]
env_conf ["env_gpu"] = gpu_id
torch.manual_seed(args.seed + rank)
if gpu_id >= 0:
torch.cuda.manual_seed(args.seed + rank)
if "EM_env" in args.env:
raw, lbl, prob, gt_lbl = datasets
env = EM_env (raw, lbl, prob, env_conf, 'train', gt_lbl)
else:
env = Voronoi_env (env_conf)
if optimizer is None:
if args.optimizer == 'RMSprop':
optimizer = optim.RMSprop (shared_model.parameters (), lr=args.lr)
if args.optimizer == 'Adam':
optimizer = optim.Adam (shared_model.parameters (), lr=args.lr, amsgrad=args.amsgrad)
# env.seed (args.seed + rank)
if not args.continuous:
player = Agent (None, env, args, None)
else:
player = Agent_continuous (None, env, args, None)
player.gpu_id = gpu_id
if not args.continuous:
player.model = A3Clstm (env.observation_space.shape, env_conf["num_action"], args.hidden_feat)
else:
player.model = A3Clstm_continuous (env.observation_space.shape, env_conf["num_action"], args.hidden_feat)
player.state = player.env.reset ()
player.state = torch.from_numpy (player.state).float ()
old_score = player.env.old_score
final_score = 0
if gpu_id >= 0:
with torch.cuda.device (gpu_id):
player.state = player.state.cuda ()
player.model = player.model.cuda ()
player.model.train ()
if rank == 0:
eps_reward = 0
pinned_eps_reward = 0
mean_log_prob = 0
# print ("rank: ", rank)
while True:
if gpu_id >= 0:
with torch.cuda.device (gpu_id):
player.model.load_state_dict (shared_model.state_dict ())
else:
player.model.load_state_dict (shared_model.state_dict ())
if player.done:
player.eps_len = 0
if rank == 0:
if 0 <= (train_step % args.train_log_period) < args.max_episode_length:
print ("train: step", train_step, "\teps_reward", eps_reward,
"\timprovement", final_score - old_score)
old_score = player.env.old_score
pinned_eps_reward = eps_reward
eps_reward = 0
mean_log_prob = 0
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.cx = Variable(torch.zeros(1, args.hidden_feat).cuda())
player.hx = Variable(torch.zeros(1, args.hidden_feat).cuda())
else:
player.cx = Variable(torch.zeros(1, args.hidden_feat))
player.hx = Variable(torch.zeros(1, args.hidden_feat))
else:
player.cx = Variable(player.cx.data)
player.hx = Variable(player.hx.data)
for step in range(args.num_steps):
player.action_train ()
if rank == 0:
# if 0 <= (train_step % args.train_log_period) < args.max_episode_length:
# print ("train: step", train_step, "\taction = ", player.action)
eps_reward += player.reward
# print (eps_reward)
mean_log_prob += player.log_probs [-1] / env_conf ["T"]
if player.done:
break
if player.done:
# if rank == 0:
# print ("----------------------------------------------")
final_score = player.env.old_score
state = player.env.reset ()
player.state = torch.from_numpy (state).float ()
if gpu_id >= 0:
with torch.cuda.device (gpu_id):
player.state = player.state.cuda ()
R = torch.zeros (1, 1)
if not player.done:
if not args.continuous:
value, _, _ = player.model((Variable(player.state.unsqueeze(0)),
(player.hx, player.cx)))
else:
value, _, _, _ = player.model((Variable(player.state.unsqueeze(0)),
(player.hx, player.cx)))
R = value.data
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
R = R.cuda()
player.values.append(Variable(R))
policy_loss = 0
value_loss = 0
gae = torch.zeros(1, 1)
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
gae = gae.cuda()
R = Variable(R)
for i in reversed(range(len(player.rewards))):
R = args.gamma * R + player.rewards[i]
advantage = R - player.values[i]
value_loss = value_loss + 0.5 * advantage.pow(2)
delta_t = player.values[i + 1].data * args.gamma + player.rewards[i] - \
player.values[i].data
gae = gae * args.gamma * args.tau + delta_t
# print (player.rewards [i])
if not args.continuous:
policy_loss = policy_loss - \
player.log_probs[i] * \
Variable(gae) - 0.01 * player.entropies[i]
else:
policy_loss = policy_loss - \
player.log_probs[i].sum () * Variable(gae) - \
0.01 * player.entropies[i].sum ()
player.model.zero_grad ()
sum_loss = (policy_loss + value_loss)
sum_loss.backward ()
ensure_shared_grads (player.model, shared_model, gpu=gpu_id >= 0)
optimizer.step ()
player.clear_actions ()
if rank == 0:
train_step += 1
if train_step % args.log_period == 0:
log_info = {
# 'train: sum_loss': sum_loss,
'train: value_loss': value_loss,
'train: policy_loss': policy_loss,
'train: advanage': advantage,
# 'train: entropy': entropy,
'train: eps reward': pinned_eps_reward,
# 'train: mean log prob': mean_log_prob
}
for tag, value in log_info.items ():
logger.scalar_summary (tag, value, train_step)
def train(rank, args, shared_model, optimizer):
ptitle('TrainingAgent{}'.format(rank))
gpu_id = args.gpu_ids[rank % len(args.gpu_ids)]
num_trained_episodes = 0
torch.manual_seed(args.seed + rank)
if gpu_id >= 0:
torch.cuda.manual_seed(args.seed + rank)
env = create_env(args.env, args)
if optimizer is None:
if args.optimizer == 'RMSprop':
optimizer = optim.RMSprop(shared_model.parameters(), lr=args.lr)
if args.optimizer == 'Adam':
optimizer = optim.Adam(shared_model.parameters(), lr=args.lr)
env.seed(args.seed + rank)
tp_weight = args.tp
player = Agent(None, env, args, None)
player.gpu_id = gpu_id
player.model = A3C_CONV(args.stack_frames, player.env.action_space, args.terminal_prediction, args.reward_prediction)
player.state = player.env.reset()
player.state = torch.from_numpy(player.state).float()
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.state = player.state.cuda()
player.model = player.model.cuda()
player.model.train()
while True:
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.model.load_state_dict(shared_model.state_dict())
else:
player.model.load_state_dict(shared_model.state_dict())
if player.done:
#sys._debugmallocstats()
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.cx = Variable(torch.zeros(1, 128).cuda())
player.hx = Variable(torch.zeros(1, 128).cuda())
else:
player.cx = Variable(torch.zeros(1, 128))
player.hx = Variable(torch.zeros(1, 128))
else:
player.cx = Variable(player.cx.data)
player.hx = Variable(player.hx.data)
for step in range(args.num_steps):
player.eps_len += 1
player.action_train()
if player.done:
num_trained_episodes += 1
break
if player.done:
state = player.env.reset()
player.state = torch.from_numpy(state).float()
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
player.state = player.state.cuda()
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
R = torch.zeros(1, 1).cuda()
else:
R = torch.zeros(1, 1)
if not player.done:
state = player.state
state = state.unsqueeze(0)
value, _, _, _, _, _ = player.model((Variable(state), (player.hx, player.cx)))
R = value.data
player.values.append(Variable(R))
policy_loss = 0
value_loss = 0
terminal_loss = 0
reward_pred_loss = 0
R = Variable(R)
if gpu_id >= 0:
with torch.cuda.device(gpu_id):
gae = torch.zeros(1, 1).cuda()
else:
gae = torch.zeros(1, 1)
for i in reversed(range(len(player.rewards))):
R = args.gamma * R + player.rewards[i]
advantage = R - player.values[i]
value_loss = value_loss + 0.5 * advantage.pow(2)
# Generalized Advantage Estimataion
# print(player.rewards[i])
delta_t = player.rewards[i] + args.gamma * \
player.values[i + 1].data - player.values[i].data
gae = gae * args.gamma * args.tau + delta_t
policy_loss = policy_loss - (player.log_probs[i].sum() * Variable(gae)) - (0.01 * player.entropies[i].sum())
if args.reward_prediction:
reward_pred_loss = reward_pred_loss + (player.reward_predictions[i] - player.rewards[i]).pow(2)
if args.terminal_prediction: # new way of using emprical episode length as a proxy for current length.
if player.average_episode_length is None:
end_predict_labels = np.arange(player.eps_len-len(player.terminal_predictions), player.eps_len) / player.eps_len # heuristic
else:
end_predict_labels = np.arange(player.eps_len-len(player.terminal_predictions), player.eps_len) / player.average_episode_length
for i in range(len(player.terminal_predictions)):
terminal_loss = terminal_loss + (player.terminal_predictions[i] - end_predict_labels[i]).pow(2)
terminal_loss /= len(player.terminal_predictions)
player.model.zero_grad()
total_loss = policy_loss + 0.5 * value_loss + tp_weight*terminal_loss + 0.5*reward_pred_loss
# Visualize Computation Graph
#graph = make_dot(total_loss)
#from graphviz import Source
#Source.view(graph)
total_loss.backward()
ensure_shared_grads(player.model, shared_model, gpu=gpu_id >= 0)
optimizer.step()
player.clear_actions()
if player.done:
#print(f' CPU {rank} -> train episode count is {num_trained_episodes}')
if player.average_episode_length is None: # initial one
player.average_episode_length = player.eps_len
else:
player.average_episode_length = int(0.99 * player.average_episode_length + 0.01 * player.eps_len)
#print(player.average_episode_length, 'current one is ', player.eps_len)
player.eps_len = 0 # reset here