def train(rank, args, shared_model, opt_ac, can_save, shared_obs_stats):
best_result = -1000
torch.manual_seed(args.seed + rank)
torch.set_default_tensor_type('torch.DoubleTensor')
num_inputs = args.feature
num_actions = 9
last_state = [1] * 48
if args.render and can_save:
env = RunEnv(visualize=True)
else:
env = RunEnv(visualize=False)
#running_state = ZFilter((num_inputs,), clip=5)
#running_reward = ZFilter((1,), demean=False, clip=10)
episode_lengths = []
PATH_TO_MODEL = '../models/' + str(args.bh)
ac_net = ActorCritic(num_inputs, num_actions)
start_time = time.time()
for i_episode in count(1):
memory = Memory()
ac_net.load_state_dict(shared_model.state_dict())
ac_net.zero_grad()
num_steps = 0
reward_batch = 0
num_episodes = 0
#Tot_loss = 0
#Tot_num =
while num_steps < args.batch_size:
#state = env.reset()
#print(num_steps)
state = env.reset(difficulty=0)
last_state = process_observation(state)
state = process_observation(state)
last_state, state = transform_observation(last_state, state)
state = numpy.array(state)
#global last_state
#last_state,_ = update_observation(last_state,state)
#last_state,state = update_observation(last_state,state)
#print(state.shape[0])
#print(state[41])
state = Variable(torch.Tensor(state).unsqueeze(0))
shared_obs_stats.observes(state)
state = shared_obs_stats.normalize(state)
state = state.data[0].numpy()
#state = running_state(state)
reward_sum = 0
#timer = time.time()
for t in range(10000): # Don't infinite loop while learning
#print(t)
if args.use_sep_pol_val:
action = select_action(state)
else:
action = select_action_actor_critic(state, ac_net)
#print(action)
action = action.data[0].numpy()
if numpy.any(numpy.isnan(action)):
print(state)
print(action)
print('ERROR')
raise RuntimeError('action NaN problem')
#print(action)
#print("------------------------")
#timer = time.time()
BB = numpy.append(action, action)
#print(BB)
reward = 0
if args.skip:
#env.step(action)
_, A, _, _ = env.step(BB)
reward += A
_, A, _, _ = env.step(BB)
reward += A
next_state, A, done, _ = env.step(BB)
reward += A
next_state = process_observation(next_state)
last_state, next_state = transform_observation(
last_state, next_state)
next_state = numpy.array(next_state)
reward_sum += reward
#print('env:')
#print(time.time()-timer)
#last_state ,next_state = update_observation(last_state,next_state)
#next_state = running_state(next_state)
next_state = Variable(torch.Tensor(next_state).unsqueeze(0))
shared_obs_stats.observes(next_state)
next_state = shared_obs_stats.normalize(next_state)
next_state = next_state.data[0].numpy()
#print(next_state[41:82])
mask = 1
if done:
mask = 0
memory.push(state, np.array([action]), mask, next_state,
reward)
#if args.render:
# env.render()
if done:
break
state = next_state
num_steps += (t - 1)
num_episodes += 1
reward_batch += reward_sum
reward_batch /= num_episodes
batch = memory.sample()
#print('env:')
#print(time.time()-timer)
#timer = time.time()
update_params_actor_critic(batch, args, shared_model, ac_net, opt_ac)
#print('backpropagate:')
#print(time.time()-timer)
epoch = i_episode
if (i_episode % args.log_interval == 0) and (rank == 0):
print('TrainEpisode {}\tLast reward: {}\tAverage reward {:.2f}'.
format(i_episode, reward_sum, reward_batch))
if reward_batch > best_result:
best_result = reward_batch
save_model(
{
'epoch': epoch,
'bh': args.bh,
'state_dict': ac_net.state_dict(),
'optimizer': opt_ac,
'obs': shared_obs_stats,
}, PATH_TO_MODEL, 'best')
if epoch % 30 == 1:
save_model(
{
'epoch': epoch,
'bh': args.bh,
'state_dict': ac_net.state_dict(),
'optimizer': opt_ac,
'obs': shared_obs_stats,
}, PATH_TO_MODEL, epoch)
def test(rank, args, shared_model, opt_ac):
best_result = -1000
torch.manual_seed(args.seed + rank)
torch.set_default_tensor_type('torch.DoubleTensor')
num_inputs = args.feature
num_actions = 9
last_state = numpy.zeros(41)
if args.render:
env = RunEnv(visualize=True)
else:
env = RunEnv(visualize=False)
running_state = ZFilter((num_inputs, ), clip=5)
running_reward = ZFilter((1, ), demean=False, clip=10)
episode_lengths = []
PATH_TO_MODEL = '../models/' + str(args.bh)
ac_net = ActorCritic(num_inputs, num_actions)
start_time = time.time()
for i_episode in count(1):
memory = Memory()
ac_net.load_state_dict(shared_model.state_dict())
num_steps = 0
reward_batch = 0
num_episodes = 0
while num_steps < args.batch_size:
#state = env.reset()
#print(num_steps)
state = env.reset(difficulty=0)
state = numpy.array(state)
#global last_state
#last_state = state
#last_state,_ = update_observation(last_state,state)
#last_state,state = update_observation(last_state,state)
#print(state.shape[0])
#print(state[41])
state = running_state(state)
reward_sum = 0
for t in range(10000): # Don't infinite loop while learning
#print(t)
#timer = time.time()
if args.use_sep_pol_val:
action = select_action(state)
else:
action = select_action_actor_critic(state, ac_net)
#print(action)
action = action.data[0].numpy()
if numpy.any(numpy.isnan(action)):
print(action)
puts('ERROR')
return
#print('NN take:')
#print(time.time()-timer)
#print(action)
#print("------------------------")
#timer = time.time()
if args.skip:
#env.step(action)
_, reward, _, _ = env.step(action)
reward_sum += reward
next_state, reward, done, _ = env.step(action)
next_state = numpy.array(next_state)
reward_sum += reward
#print('env take:')
#print(time.time()-timer)
#timer = time.time()
#last_state ,next_state = update_observation(last_state,next_state)
next_state = running_state(next_state)
#print(next_state[41:82])
mask = 1
if done:
mask = 0
#print('update take:')
#print(time.time()-timer)
#timer = time.time()
memory.push(state, np.array([action]), mask, next_state,
reward)
#print('memory take:')
#print(time.time()-timer)
#if args.render:
# env.render()
if done:
break
state = next_state
num_steps += (t - 1)
num_episodes += 1
#print(num_episodes)
reward_batch += reward_sum
#print(num_episodes)
reward_batch /= num_episodes
batch = memory.sample()
#update_params_actor_critic(batch,args,shared_model,ac_net,opt_ac)
time.sleep(60)
if i_episode % args.log_interval == 0:
File = open(PATH_TO_MODEL + '/record.txt', 'a+')
File.write("Time {}, episode reward {}, Average reward {}".format(
time.strftime("%Hh %Mm %Ss",
time.gmtime(time.time() - start_time)),
reward_sum, reward_batch))
File.close()
#print('TestEpisode {}\tLast reward: {}\tAverage reward {:.2f}'.format(
# i_episode, reward_sum, reward_batch))
print("Time {}, episode reward {}, Average reward {}".format(
time.strftime("%Hh %Mm %Ss",
time.gmtime(time.time() - start_time)),
reward_sum, reward_batch))
#print('!!!!')
epoch = i_episode
if reward_batch > best_result:
best_result = reward_batch
save_model(
{
'epoch': epoch,
'bh': args.bh,
'state_dict': shared_model.state_dict(),
'optimizer': opt_ac.state_dict(),
}, PATH_TO_MODEL, 'best')
if epoch % 30 == 1:
save_model(
{
'epoch': epoch,
'bh': args.bh,
'state_dict': shared_model.state_dict(),
'optimizer': opt_ac.state_dict(),
}, PATH_TO_MODEL, epoch)
traffic_light = TrafficLight()
counter = Counter()
ac_net = ActorCritic(num_inputs, num_actions)
opt_ac = optim.Adam(ac_net.parameters(), lr=args.lr)
shared_grad_buffers = Shared_grad_buffers(ac_net)
shared_obs_stats = Shared_obs_stats(num_inputs)
if args.resume:
print("=> loading checkpoint ")
checkpoint = torch.load('../../7.87.t7')
#checkpoint = torch.load('../../best.t7')
args.start_epoch = checkpoint['epoch']
#best_prec1 = checkpoint['best_prec1']
ac_net.load_state_dict(checkpoint['state_dict'])
opt_ac.load_state_dict(checkpoint['optimizer'])
opt_ac.state = defaultdict(dict, opt_ac.state)
#print(opt_ac)
shared_obs_stats = checkpoint['obs']
print(ac_net)
print("=> loaded checkpoint (epoch {})".format(checkpoint['epoch']))
ac_net.share_memory()
#opt_ac.share_memory()
#running_state = ZFilter((num_inputs,), clip=5)
processes = []
def train(rank, args, traffic_light, counter, shared_model,
shared_grad_buffers, shared_obs_stats, opt_ac):
best_result = -1000
torch.manual_seed(args.seed + rank)
torch.set_default_tensor_type('torch.DoubleTensor')
num_inputs = args.feature
num_actions = 9
last_state = [0] * 41
last_v = [0] * 10
#last_state = numpy.zeros(48)
env = RunEnv(visualize=False)
#running_state = ZFilter((num_inputs,), clip=5)
#running_reward = ZFilter((1,), demean=False, clip=10)
episode_lengths = []
PATH_TO_MODEL = '../models/' + str(args.bh)
ac_net = ActorCritic(num_inputs, num_actions)
#running_state = ZFilter((num_inputs,), clip=5)
start_time = time.time()
for i_episode in range(args.start_epoch + 1, 999999):
#print(shared_obs_stats.n[0])
#print('hei')
#if rank == 0:
# print(running_state.rs._n)
signal_init = traffic_light.get()
memory = Memory()
ac_net.load_state_dict(shared_model.state_dict())
num_steps = 0
reward_batch = 0
num_episodes = 0
#Tot_loss = 0
#Tot_num =
while num_steps < args.batch_size:
#state = env.reset()
#print(num_steps)
state = env.reset(difficulty=0)
#state = numpy.array(state)
last_state, last_v, state = process_observation(
last_state, last_v, state)
state = numpy.array(state)
#state = running_state(state)
state = Variable(torch.Tensor(state).unsqueeze(0))
shared_obs_stats.observes(state)
state = shared_obs_stats.normalize(state)
state = state.data[0].numpy()
#print(state)
#return
#print(AA)
#print(type(AA))
#print(type(state))
#print(AA.shape)
#print(state.shape)
reward_sum = 0
#timer = time.time()
for t in range(10000): # Don't infinite loop while learning
#print(t)
if args.use_sep_pol_val:
action = select_action(state)
else:
action = select_action_actor_critic(state, ac_net)
#print(action)
action = action.data[0].numpy()
if numpy.any(numpy.isnan(action)):
print(state)
print(action)
print(ac_net.affine1.weight)
print(ac_net.affine1.weight.data)
print('ERROR')
#action = select_action_actor_critic(state,ac_net)
#action = action.data[0].numpy()
#state = state + numpy.random.rand(args.feature)*0.001
raise RuntimeError('action NaN problem')
#print(action)
#print("------------------------")
#timer = time.time()
reward = 0
if args.skip:
#env.step(action)
_, A, _, _ = env.step(action)
reward += A
_, A, _, _ = env.step(action)
reward += A
BB = numpy.append(action, action)
next_state, A, done, _ = env.step(BB)
reward += A
#print(next_state)
#last_state = process_observation(state)
last_state, last_v, next_state = process_observation(
last_state, last_v, next_state)
next_state = numpy.array(next_state)
#print(next_state)
#print(next_state.shape)
#return
reward_sum += reward
#print('env:')
#print(time.time()-timer)
#last_state ,next_state = update_observation(last_state,next_state)
#next_state = running_state(next_state)
next_state = Variable(torch.Tensor(next_state).unsqueeze(0))
shared_obs_stats.observes(next_state)
next_state = shared_obs_stats.normalize(next_state)
next_state = next_state.data[0].numpy()
#print(next_state[41:82])
mask = 1
if done:
mask = 0
memory.push(state, np.array([action]), mask, next_state,
reward)
#if args.render:
# env.render()
if done:
break
state = next_state
num_steps += (t - 1)
num_episodes += 1
reward_batch += reward_sum
reward_batch /= num_episodes
batch = memory.sample()
#print('env:')
#print(time.time()-timer)
#timer = time.time()
update_params_actor_critic(batch, args, ac_net, opt_ac)
shared_grad_buffers.add_gradient(ac_net)
counter.increment()
epoch = i_episode
if (i_episode % args.log_interval == 0) and (rank == 0):
print(
'TrainEpisode {}\tTime{}\tLast reward: {}\tAverage reward {:.2f}'
.format(
i_episode,
time.strftime("%Hh %Mm %Ss",
time.gmtime(time.time() - start_time)),
reward_sum, reward_batch))
epoch = i_episode
if reward_batch > best_result:
best_result = reward_batch
save_model(
{
'epoch': epoch,
'bh': args.bh,
'state_dict': shared_model.state_dict(),
'optimizer': opt_ac.state_dict(),
'obs': shared_obs_stats,
}, PATH_TO_MODEL, 'best')
if epoch % 30 == 1:
save_model(
{
'epoch': epoch,
'bh': args.bh,
'state_dict': shared_model.state_dict(),
'optimizer': opt_ac.state_dict(),
'obs': shared_obs_stats,
}, PATH_TO_MODEL, epoch)
# wait for a new signal to continue
while traffic_light.get() == signal_init:
pass
def play(args):
env = create_mario_env(args.env_name, ACTIONS[args.move_set])
observation_space = env.observation_space.shape[0]
action_space = env.action_space.n
model = ActorCritic(observation_space, action_space)
checkpoint_file = \
f"{args.env_name}/{args.model_id}_{args.algorithm}_params.tar"
checkpoint = restore_checkpoint(checkpoint_file)
assert args.env_name == checkpoint['env'], \
"This checkpoint is for different environment: {checkpoint['env']}"
args.model_id = checkpoint['id']
print(f"Environment: {args.env_name}")
print(f" Agent: {args.model_id}")
model.load_state_dict(checkpoint['model_state_dict'])
state = env.reset()
state = torch.from_numpy(state)
reward_sum = 0
done = True
episode_length = 0
start_time = time.time()
for step in count():
episode_length += 1
# shared model sync
if done:
cx = torch.zeros(1, 512)
hx = torch.zeros(1, 512)
else:
cx = cx.data
hx = hx.data
with torch.no_grad():
value, logit, (hx, cx) = model((state.unsqueeze(0), (hx, cx)))
prob = F.softmax(logit, dim=-1)
action = prob.max(-1, keepdim=True)[1]
action_idx = action.item()
action_out = ACTIONS[args.move_set][action_idx]
state, reward, done, info = env.step(action_idx)
reward_sum += reward
print(
f"{emojize(':mushroom:')} World {info['world']}-{info['stage']} | {emojize(':video_game:')}: [ {' + '.join(action_out):^13s} ] | ",
end='\r',
)
env.render()
if done:
t = time.time() - start_time
print(
f"{emojize(':mushroom:')} World {info['world']}-{info['stage']} |" + \
f" {emojize(':video_game:')}: [ {' + '.join(action_out):^13s} ] | " + \
f"ID: {args.model_id}, " + \
f"Time: {time.strftime('%H:%M:%S', time.gmtime(t)):^9s}, " + \
f"Reward: {reward_sum: 10.2f}, " + \
f"Progress: {(info['x_pos'] / 3225) * 100: 3.2f}%",
end='\r',
flush=True,
)
reward_sum = 0
episode_length = 0
time.sleep(args.reset_delay)
state = env.reset()
state = torch.from_numpy(state)
offset = 20
width = window.width - offset
height = window.height - offset
board_unit = min(width // board_size, height // board_size)
x1_board = window.width // 2 - (board_size // 2 + 1) * board_unit
x2_board = x1_board + (board_size + 1) * board_unit
y1_board = window.height // 2 - (board_size // 2 + 1) * board_unit
y2_board = y1_board + (board_size + 1) * board_unit
print(x1_board, x2_board, y1_board, y2_board)
env = TrainEnvSingle()
game = env.game
model = ActorCritic()
model.load_state_dict(torch.load("weights.pt"))
model = model.eval()
state, invalid = env.reset()
dist, value = model(state, invalid)
q_values = dist.probs.tolist()[0]
def take_action(dt):
pass
def reload_model(dt):
global model
model.load_state_dict(torch.load("weights.pt"))
print("Reloaded model")
class ActorCriticAgentUsingICM:
def __init__(self, nb_actions, learning_rate, gamma, hidden_size,
model_input_size, entropy_coeff_start, entropy_coeff_end,
entropy_coeff_anneal, continuous):
self.device = torch.device(
"cuda" if torch.cuda.is_available() else "cpu")
self.num_actions = nb_actions
self.gamma = gamma
self.continuous = continuous
self.learning_rate = learning_rate
self.entropy_coefficient_start = entropy_coeff_start
self.entropy_coefficient_end = entropy_coeff_end
self.entropy_coefficient_anneal = entropy_coeff_anneal
self.step_no = 0
if self.continuous:
self.model = ActorCriticContinuous(hidden_size=hidden_size,
inputs=model_input_size,
outputs=nb_actions).to(
self.device)
else:
self.model = ActorCritic(hidden_size=hidden_size,
inputs=model_input_size,
outputs=nb_actions).to(self.device)
self.hidden_size = hidden_size
self.optimizer = torch.optim.Adam(self.model.parameters(),
lr=self.learning_rate)
self.loss_function = torch.nn.MSELoss()
self.memory = []
self.ICM = ICM(model_input_size, nb_actions)
self.ICM.train()
# Get the current entropy coefficient value according to the start/end and annealing values
def get_entropy_coefficient(self):
entropy = self.entropy_coefficient_end
if self.step_no < self.entropy_coefficient_anneal:
entropy = self.entropy_coefficient_start - self.step_no * \
((self.entropy_coefficient_start - self.entropy_coefficient_end) /
self.entropy_coefficient_anneal)
return entropy
# select an action with policy
def select_action(self, state):
self.step_no += 1
if self.continuous:
action_mean, action_dev, state_value = self.model(state)
action_dist = Normal(action_mean, action_dev)
else:
action_probs, state_value = self.model(state)
action_dist = Categorical(action_probs)
return action_dist, state_value
def update_model(self):
Gt = torch.tensor(0)
policy_losses = []
forward_losses = []
inverse_losses = []
value_losses = []
entropy_loss = []
returns = []
# calculate the true value using rewards returned from the environment
for (_, reward, _, _, _, _, _) in self.memory[::-1]:
# calculate the discounted value
Gt = reward + self.gamma * Gt
returns.insert(0, Gt)
returns = torch.tensor(returns)
returns = (returns - returns.mean()) / (returns.std() + 1e-8)
for (action_prob, _, state_value, entropy, state, next_state,
action), Gt in zip(self.memory, returns):
advantage = Gt.item() - state_value.item()
# calculate actor (policy) loss
policy_losses.append((-action_prob * advantage).mean())
# calculate critic (value) loss using model loss function
value_losses.append(
self.loss_function(state_value, Gt.unsqueeze(0)))
entropy_loss.append(-entropy)
forward_losses.append(
self.ICM.get_forward_loss(state, action, next_state))
inverse_losses.append(
self.ICM.get_inverse_loss(state, action, next_state))
# reset gradients
self.optimizer.zero_grad()
self.ICM.optimizer.zero_grad()
# sum up all the values of policy_losses and value_losses
icm_loss = (1 - self.ICM.beta) * torch.stack(inverse_losses).mean(
) + self.ICM.beta * torch.stack(forward_losses).mean()
loss = self.ICM.lambda_weight*(torch.stack(policy_losses).mean() + \
torch.stack(value_losses).mean() + self.get_entropy_coefficient() * \
torch.stack(entropy_loss).mean()) + icm_loss
loss.backward()
self.optimizer.step()
self.ICM.optimizer.step()
self.memory = []
return loss.item()
# save model
def save(self, path, name):
dirname = os.path.dirname(__file__)
filename = os.path.join(dirname, os.path.join(path, name + ".pt"))
torch.save(self.model.state_dict(), filename)
# load a model
def load(self, path):
dirname = os.path.dirname(__file__)
filename = os.path.join(dirname, path)
self.model.load_state_dict(torch.load(filename))
def cache(self, action_prob, reward, state_value, entropy, state,
next_state, action):
self.memory.append((action_prob, reward, state_value, entropy, state,
next_state, action))
def test(rank, args, shared_model, counter, device):
# time.sleep(10.)
# logging
log_dir = f'logs/{args.env_name}/{args.model_id}/{args.uuid}/'
info_logger = setup_logger('info', log_dir, f'info.log')
result_logger = setup_logger('results', log_dir, f'results.log')
# torch.manual_seed(args.seed + rank)
env = create_atari_environment(args.env_name)
if args.record:
if not os.path.exists(f'playback/{args.env_name}/'):
os.makedirs(f'playback/{args.env_name}/{args.model_id}', exist_ok=True)
env = gym.wrappers.Monitor(env, f'playback/{args.env_name}/{args.model_id}/', force=True)
# env.seed(args.seed + rank)
observation_space = env.observation_space.shape[0]
action_space = env.action_space.n
model = ActorCritic(observation_space, action_space)
if torch.cuda.is_available():
model.cuda()
model.eval()
state = env.reset()
state = torch.from_numpy(state)
reward_sum = 0
done = True
episode_length = 0
actions = deque(maxlen=4000)
start_time = time.time()
for episode in count():
episode_length += 1
# shared model sync
if done:
model.load_state_dict(shared_model.state_dict())
cx = torch.zeros(1, 512)
hx = torch.zeros(1, 512)
else:
cx = cx.data
hx = hx.data
with torch.no_grad():
value, logit, (hx, cx) = model((state.unsqueeze(0), (hx, cx)))
prob = F.softmax(logit, dim=-1)
action = prob.max(-1, keepdim=True)[1]
state, reward, done, info = env.step(action.item())
reward_sum += reward
info_log = {
'id': args.model_id,
'algorithm': args.algorithm,
'greedy-eps': args.greedy_eps,
'episode': episode,
'total_episodes': counter.value,
'episode_length': episode_length,
'reward': reward_sum,
'done': done,
}
info_logger.info(info_log)
print(f"{emojize(':video_game:', use_aliases=True)} | ", end='\r')
env.render()
actions.append(action.item())
if done:
t = time.time() - start_time
print(
f"{emojize(':video_game:', use_aliases=True)} | " + \
f"ID: {args.model_id}, " + \
f"Total Episodes: {counter.value}, " + \
f"Time: {time.strftime('%H:%M:%S', time.gmtime(t)):^9s}, " + \
f"FPS: {episode_length/t: 6.2f}, " + \
f"Reward: {reward_sum: 10.0f}",
end='\r',
flush=True,
)
result_logger.info(info_log)
reward_sum = 0
episode_length = 0
actions.clear()
time.sleep(args.reset_delay)
state = env.reset()
state = torch.from_numpy(state)
# create net
action_size = env.action_space.shape[0]
number_asset, seq_window, features_all = env.observation_space.shape
assert action_size == number_asset + 1
input_size = features_all - 1
net = ActorCritic(input_size=input_size,
hidden_size=50,
action_size=action_size)
net_tgt = ActorCritic(input_size=input_size,
hidden_size=50,
action_size=action_size)
net_tgt.eval()
print(net_tgt)
net_tgt.load_state_dict(net.state_dict())
# create replay
replay = ch.ExperienceReplay()
# create loss function
criterion_mse = nn.MSELoss()
# create optimizer
optimizer_actor = torch.optim.Adam(net.actor.parameters(), lr=0.001)
optimizer_critic = torch.optim.Adam(net.critic.parameters(), lr=0.001)
def update(replay):
# batch-data
state_batch = replay.state()
def main(args):
print(f" Session ID: {args.uuid}")
# logging
log_dir = f'logs/{args.env_name}/{args.model_id}/{args.uuid}/'
args_logger = setup_logger('args', log_dir, f'args.log')
env_logger = setup_logger('env', log_dir, f'env.log')
if args.debug:
debug.packages()
os.environ['OMP_NUM_THREADS'] = "1"
if torch.cuda.is_available():
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
devices = ",".join([str(i) for i in range(torch.cuda.device_count())])
os.environ["CUDA_VISIBLE_DEVICES"] = devices
args_logger.info(vars(args))
env_logger.info(vars(os.environ))
env = create_atari_environment(args.env_name)
shared_model = ActorCritic(env.observation_space.shape[0],
env.action_space.n)
if torch.cuda.is_available():
shared_model = shared_model.cuda()
shared_model.share_memory()
optimizer = SharedAdam(shared_model.parameters(), lr=args.lr)
optimizer.share_memory()
if args.load_model: # TODO Load model before initializing optimizer
checkpoint_file = f"{args.env_name}/{args.model_id}_{args.algorithm}_params.tar"
checkpoint = restore_checkpoint(checkpoint_file)
assert args.env_name == checkpoint['env'], \
"Checkpoint is for different environment"
args.model_id = checkpoint['id']
args.start_step = checkpoint['step']
print("Loading model from checkpoint...")
print(f"Environment: {args.env_name}")
print(f" Agent: {args.model_id}")
print(f" Start: Step {args.start_step}")
shared_model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
else:
print(f"Environment: {args.env_name}")
print(f" Agent: {args.model_id}")
torch.manual_seed(args.seed)
print(
FontColor.BLUE + \
f"CPUs: {mp.cpu_count(): 3d} | " + \
f"GPUs: {None if not torch.cuda.is_available() else torch.cuda.device_count()}" + \
FontColor.END
)
processes = []
counter = mp.Value('i', 0)
lock = mp.Lock()
# Queue training processes
num_processes = args.num_processes
no_sample = args.non_sample # count of non-sampling processes
if args.num_processes > 1:
num_processes = args.num_processes - 1
samplers = num_processes - no_sample
for rank in range(0, num_processes):
device = 'cpu'
if torch.cuda.is_available():
device = 0 # TODO: Need to move to distributed to handle multigpu
if rank < samplers: # random action
p = mp.Process(
target=train,
args=(rank, args, shared_model, counter, lock, optimizer,
device),
)
else: # best action
p = mp.Process(
target=train,
args=(rank, args, shared_model, counter, lock, optimizer,
device, False),
)
p.start()
time.sleep(1.)
processes.append(p)
# Queue test process
p = mp.Process(target=test,
args=(args.num_processes, args, shared_model, counter, 0))
p.start()
processes.append(p)
for p in processes:
p.join()
def train(rank,
args,
shared_model,
counter,
lock,
optimizer=None,
device='cpu',
select_sample=True):
# torch.manual_seed(args.seed + rank)
# logging
log_dir = f'logs/{args.env_name}/{args.model_id}/{args.uuid}/'
loss_logger = setup_logger('loss', log_dir, f'loss.log')
# action_logger = setup_logger('actions', log_dir, f'actions.log')
text_color = FontColor.RED if select_sample else FontColor.GREEN
print(
text_color +
f"Process: {rank: 3d} | {'Sampling' if select_sample else 'Decision'} | Device: {str(device).upper()}",
FontColor.END)
env = create_atari_environment(args.env_name)
observation_space = env.observation_space.shape[0]
action_space = env.action_space.n
# env.seed(args.seed + rank)
model = ActorCritic(observation_space, action_space)
if torch.cuda.is_available():
model = model.cuda()
model.device = device
if optimizer is None:
optimizer = optim.Adam(shared_model.parameters(), lr=args.lr)
model.train()
state = env.reset()
state = torch.from_numpy(state)
done = True
for t in count(start=args.start_step):
if t % args.save_interval == 0 and t > 0:
save_checkpoint(shared_model, optimizer, args, t)
# Sync shared model
model.load_state_dict(shared_model.state_dict())
if done:
cx = torch.zeros(1, 512)
hx = torch.zeros(1, 512)
else:
cx = cx.detach()
hx = hx.detach()
values = []
log_probs = []
rewards = []
entropies = []
episode_length = 0
for step in range(args.num_steps):
episode_length += 1
value, logit, (hx, cx) = model((state.unsqueeze(0), (hx, cx)))
prob = F.softmax(logit, dim=-1)
log_prob = F.log_softmax(logit, dim=-1)
entropy = -(log_prob * prob).sum(-1, keepdim=True)
entropies.append(entropy)
reason = ''
if select_sample:
rand = random.random()
epsilon = get_epsilon(t)
if rand < epsilon and args.greedy_eps:
action = torch.randint(0, action_space, (1, 1))
reason = 'uniform'
else:
action = prob.multinomial(1)
reason = 'multinomial'
else:
action = prob.max(-1, keepdim=True)[1]
reason = 'choice'
# action_logger.info({
# 'rank': rank,
# 'action': action.item(),
# 'reason': reason,
# })
if torch.cuda.is_available():
action = action.cuda()
value = value.cuda()
log_prob = log_prob.gather(-1, action)
# action_out = ACTIONS[args.move_set][action.item()]
state, reward, done, info = env.step(action.item())
done = done or episode_length >= args.max_episode_length
reward = max(min(reward, 50), -50) # h/t @ArvindSoma
with lock:
counter.value += 1
if done:
episode_length = 0
state = env.reset()
state = torch.from_numpy(state)
values.append(value)
log_probs.append(log_prob)
rewards.append(reward)
if done:
break
R = torch.zeros(1, 1)
if not done:
value, _, _ = model((state.unsqueeze(0), (hx, cx)))
R = value.data
values.append(R)
loss = gae(R, rewards, values, log_probs, entropies, args)
loss_logger.info({
'episode': t,
'rank': rank,
'sampling': select_sample,
'loss': loss.item()
})
optimizer.zero_grad()
(loss).backward()
nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)
ensure_shared_grads(model, shared_model)
optimizer.step()